Refactor toolkit

src/openpecha/pecha/__init__.py

@@ -27,7 +27,7 @@ def __init__(self, pecha_id: str, pecha_path: Path) -> None:
         self.pecha_path = pecha_path
         self.metadata = self.load_metadata()
         self.bases = self.load_bases()
-        # self.annotations = self.load_annotations()
+        self.annotations = []
 
     @classmethod
     def from_path(cls, pecha_path: Path) -> "Pecha":
@@ -65,21 +65,28 @@ def create(cls, output_path: Optional[Path] = None, pecha_id: Optional[str] = No
         return cls(pecha_id, pecha_path)
 
     @classmethod
-    def create_pecha(cls, pecha_id: str, base_text: str, annotation_id: str, annotation: List[BaseAnnotation]) -> "Pecha":
+    def create_pecha(cls, pecha_id: str, base_text: str, annotation_id: str, annotation: List[BaseAnnotation], annotation_type: AnnotationType) -> "Pecha":
         pecha = cls.create(pecha_id=pecha_id)
         base_name = pecha.set_base(base_text)
-        ann_type = get_annotation_type(annotation)
-        ann_store, _ = pecha.add_layer(base_name=base_name, layer_type=ann_type, annotation_id=annotation_id)
-
+        ann_store, _ = pecha.add_layer(base_name=base_name, layer_type=annotation_type, annotation_id=annotation_id)
         for single_annotation in annotation:
-            ann_store = pecha.add_annotation(ann_store=ann_store, annotation=single_annotation, layer_type=ann_type)
+            ann_store = pecha.add_annotation(ann_store=ann_store, annotation=single_annotation, layer_type=annotation_type)
             ann_store.save()
+        annotations = get_anns(ann_store, include_span=True)
+        for annotation in annotations:
+            pecha.annotations.append({
+                "span": {
+                    "start": annotation["span"]["start"],
+                    "end": annotation["span"]["end"],
+                },
+                "id": annotation["id"]
+            })
         return pecha
 
 
-    def add(self, annotation_id: str, annotation: List[BaseAnnotation]) -> "Pecha":
+    def add(self, annotation_id: str, annotation: List[BaseAnnotation], annotation_type: AnnotationType) -> "Pecha":
         base_name = next(iter(self.bases))
-        ann_type = get_annotation_type(annotation)
+        ann_type = annotation_type
         if check_annotation_exists(self.layer_path/base_name/f"{ann_type.value}-{annotation_id}.json"):
             raise ValueError(f"Annotation with id {annotation_id} already exists")
         ann_store, _ = self.add_layer(base_name=base_name, layer_type=ann_type, annotation_id=annotation_id)
@@ -189,7 +196,6 @@ def add_annotation(
         # Add Annotation Group Type
         ann_group_type = layer_type.annotation_group_type
         ann_data[ann_group_type.value] = layer_type.value
-
         start, end = (
             annotation.span.start,
             annotation.span.end,
@@ -219,6 +225,7 @@ def add_annotation(
             raise StamAddAnnotationError(
                 f"[Error] Failed to add annotation to STAM: {e}"
             )
+
         return ann_store
 
     def set_metadata(self, pecha_metadata: Dict):

src/openpecha/pecha/annotations.py

@@ -40,7 +40,6 @@ def end_must_not_be_less_than_start(self) -> "span":
 
 class BaseAnnotation(BaseModel):
     span: span
-    metadata: Optional[Dict] = None
 
     model_config = ConfigDict(extra="allow")
 
@@ -54,11 +53,11 @@ def get_dict(self):
 
 
 class SegmentationAnnotation(BaseAnnotation):
-    index: int
+    id: str = Field(..., description="Annotation ID")
 
 
 class AlignmentAnnotation(BaseAnnotation):
-    index: int
+    id: str = Field(..., description="Annotation ID")
     alignment_index: list[int] = Field(
         description="Index of the alignment, which can be of translation or commentary"
     )

src/openpecha/pecha/parsers/edition.py

@@ -16,6 +16,7 @@
 from openpecha.pecha.layer import AnnotationType
 from openpecha.pecha.parsers import update_coords
 from openpecha.pecha.serializers.json import JsonSerializer
+from openpecha.ids import get_annotation_id
 
 logger = get_logger(__name__)
 
@@ -44,8 +45,8 @@ def parse_segmentation(self, segments: list[str]) -> list[SegmentationAnnotation
         for index, segment in enumerate(segments, start=1):
             anns.append(
                 SegmentationAnnotation(
+                    id=str(index),
                     span=span(start=char_count, end=char_count + len(segment)),
-                    index=index,
                 )
             )
             char_count += len(segment) + 1

src/openpecha/utils.py

@@ -58,5 +58,5 @@ def write_json(
 def convert_to_base_annotation(raw_annotation):
     span_data = raw_annotation["span"]
     annotation_span = span(start=span_data["start"], end=span_data["end"])
-    annotation_data = {k: v for k, v in raw_annotation.items() if k != "span"}
+    annotation_data = {k: v for k, v in raw_annotation.items() if k != "span" and k != "index"}
     return BaseAnnotation(span=annotation_span, **annotation_data)