Merge pull request #832 from superannotateai/wm_classes

Wm classes

Author: VaghinakDev

src/superannotate/lib/app/interface/sdk_interface.py

@@ -1973,6 +1973,118 @@ def search_annotation_classes(
             raise AppException(response.errors)
         return BaseSerializer.serialize_iterable(response.data)
 
+    def update_annotation_class(
+        self,
+        project: Union[NotEmptyStr, int],
+        name: NotEmptyStr,
+        attribute_groups: List[dict],
+    ):
+        """Updates an existing annotation class by submitting a full, updated attribute_groups payload.
+        You can add new attribute groups, add new attribute values, rename attribute groups, rename attribute values,
+        delete attribute groups, delete attribute values, update attribute group types, update default attributes,
+        and update the required state.
+
+        .. warning::
+            This operation replaces the entire attribute group structure of the annotation class.
+            Any attribute groups or attribute values omitted from the payload will be permanently removed.
+            Existing annotations that reference removed attribute groups or attributes will lose their associated values.
+
+        :param project: The name or ID of the project.
+        :type project: Union[str, int]
+
+        :param name: The name of the annotation class to update.
+        :type name: str
+
+        :param attribute_groups: The full list of attribute groups for the class.
+            Each attribute group may contain:
+
+            - id (optional, required for existing groups)
+            - group_type (required): "radio", "checklist", "text", "numeric", or "ocr"
+            - name (required)
+            - isRequired (optional)
+            - default_value (optional)
+            - attributes (required, list)
+
+            Each attribute may contain:
+
+            - id (optional, required for existing attributes)
+            - name (required)
+
+        :type attribute_groups: list of dicts
+
+        Request Example:
+        ::
+
+            # Retrieve existing annotation class
+            classes = client.search_annotation_classes(project="Medical Project", name="Organ")
+            existing_class = classes[0]
+
+            # Modify attribute groups
+            updated_groups = existing_class["attribute_groups"]
+
+            # Add a new attribute to an existing group
+            updated_groups[0]["attributes"].append({"name": "Kidney"})
+
+            # Add a new attribute group
+            updated_groups.append({
+                "group_type": "radio",
+                "name": "Severity",
+                "attributes": [
+                    {"name": "Mild"},
+                    {"name": "Moderate"},
+                    {"name": "Severe"}
+                ],
+                "default_value": "Mild"
+            })
+
+            # Update the annotation class
+            client.update_annotation_class(
+                project="Medical Project",
+                name="Organ",
+                attribute_groups=updated_groups
+            )
+        """
+        project = self.controller.get_project(project)
+
+        # Find the annotation class by nam
+        annotation_classes = self.controller.annotation_classes.list(
+            condition=Condition("project_id", project.id, EQ)
+        ).data
+
+        annotation_class = next(
+            (c for c in annotation_classes if c["name"] == name), None
+        )
+
+        if not annotation_class:
+            raise AppException("Annotation class not found in project.")
+
+        # Parse and validate attribute groups
+        annotation_class["attributeGroups"] = attribute_groups
+        for group in annotation_class["attributeGroups"]:
+            if "isRequired" in group:
+                group["is_required"] = group.pop("isRequired")
+
+        from lib.core.entities import WMAnnotationClassEntity
+
+        try:
+            # validate annotation class
+            annotation_class = WMAnnotationClassEntity.parse_obj(
+                BaseSerializer(annotation_class).serialize()
+            )
+        except ValidationError as e:
+            raise AppException(wrap_error(e))
+
+        # Update the annotation class with new attribute groups
+
+        response = self.controller.annotation_classes.update(
+            project=project, annotation_class=annotation_class
+        )
+
+        if response.errors:
+            raise AppException(response.errors)
+
+        return BaseSerializer(response.data).serialize(by_alias=True)
+
     def set_project_status(self, project: NotEmptyStr, status: PROJECT_STATUS):
         """Set project status
 
@@ -2845,7 +2957,9 @@ def create_annotation_class(
         )
         if response.errors:
             raise AppException(response.errors)
-        return BaseSerializer(response.data).serialize(exclude_unset=True)
+        return BaseSerializer(response.data).serialize(
+            exclude_unset=True, by_alias=False
+        )
 
     def delete_annotation_class(
         self, project: NotEmptyStr, annotation_class: Union[dict, NotEmptyStr]

src/superannotate/lib/core/entities/__init__.py

@@ -24,8 +24,10 @@
 from lib.core.entities.project import WorkflowEntity
 from lib.core.entities.project_entities import BaseEntity
 from lib.core.entities.project_entities import S3FileEntity
+from lib.core.entities.work_managament import WMAnnotationClassEntity
 from lib.core.entities.work_managament import WMProjectUserEntity
 
+
 __all__ = [
     # base
     "ConfigEntity",
@@ -60,4 +62,6 @@
     # multimodal
     "FormModel",
     "generate_classes_from_form",
+    # work management
+    "WMAnnotationClassEntity",
 ]

src/superannotate/lib/core/entities/base.py

@@ -20,10 +20,23 @@
 try:
     from pydantic import AbstractSetIntStr  # noqa
     from pydantic import MappingIntStrAny  # noqa
+
 except ImportError:
     pass
 _missing = object()
 
+from lib.core.pydantic_v1 import Color
+from lib.core.pydantic_v1 import ColorType
+from lib.core.pydantic_v1 import validator
+
+
+class HexColor(BaseModel):
+    __root__: ColorType
+
+    @validator("__root__", pre=True)
+    def validate_color(cls, v):
+        return "#{:02X}{:02X}{:02X}".format(*Color(v).as_rgb_tuple())
+
 
 class StringDate(datetime):
     @classmethod

src/superannotate/lib/core/entities/classes.py

@@ -1,35 +1,24 @@
-from datetime import datetime
 from enum import Enum
 from typing import Any
 from typing import List
 from typing import Optional
 
-from lib.core.entities.base import BaseModel
-from lib.core.entities.base import parse_datetime
+from lib.core.entities.base import HexColor
+from lib.core.entities.base import TimedBaseModel
 from lib.core.enums import BaseTitledEnum
 from lib.core.enums import ClassTypeEnum
-from lib.core.pydantic_v1 import Color
-from lib.core.pydantic_v1 import ColorType
 from lib.core.pydantic_v1 import Extra
 from lib.core.pydantic_v1 import Field
 from lib.core.pydantic_v1 import StrictInt
 from lib.core.pydantic_v1 import StrictStr
-from lib.core.pydantic_v1 import validator
+
 
 DATE_REGEX = r"\d{4}-[01]\d-[0-3]\dT[0-2]\d:[0-5]\d:[0-5]\d(?:\.\d{3})Z"
 DATE_TIME_FORMAT_ERROR_MESSAGE = (
     "does not match expected format YYYY-MM-DDTHH:MM:SS.fffZ"
 )
 
 
-class HexColor(BaseModel):
-    __root__: ColorType
-
-    @validator("__root__", pre=True)
-    def validate_color(cls, v):
-        return "#{:02X}{:02X}{:02X}".format(*Color(v).as_rgb_tuple())
-
-
 class GroupTypeEnum(str, Enum):
     RADIO = "radio"
     CHECKLIST = "checklist"
@@ -38,22 +27,6 @@ class GroupTypeEnum(str, Enum):
     OCR = "ocr"
 
 
-class StringDate(datetime):
-    @classmethod
-    def __get_validators__(cls):
-        yield parse_datetime
-        yield cls.validate
-
-    @classmethod
-    def validate(cls, v: datetime):
-        return v.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
-
-
-class TimedBaseModel(BaseModel):
-    createdAt: StringDate = Field(None, alias="createdAt")
-    updatedAt: StringDate = Field(None, alias="updatedAt")
-
-
 class Attribute(TimedBaseModel):
     id: Optional[StrictInt]
     group_id: Optional[StrictInt]

src/superannotate/lib/core/entities/work_managament.py

@@ -6,14 +6,19 @@
 from typing import Optional
 from typing import Union
 
+from lib.core.entities.base import HexColor
 from lib.core.entities.base import TimedBaseModel
+from lib.core.enums import WMClassTypeEnum
+from lib.core.enums import WMGroupTypeEnum
 from lib.core.enums import WMUserStateEnum
 from lib.core.exceptions import AppException
 from lib.core.pydantic_v1 import BaseModel
 from lib.core.pydantic_v1 import Extra
 from lib.core.pydantic_v1 import Field
 from lib.core.pydantic_v1 import parse_datetime
 from lib.core.pydantic_v1 import root_validator
+from lib.core.pydantic_v1 import StrictInt
+from lib.core.pydantic_v1 import StrictStr
 from lib.core.pydantic_v1 import validator
 
 
@@ -209,3 +214,71 @@ def check_weight_and_value(cls, values):
         ):
             raise AppException("Weight and Value must both be set or both be None.")
         return values
+
+
+class WMAttribute(TimedBaseModel):
+    id: Optional[StrictInt]
+    group_id: Optional[StrictInt]
+    project_id: Optional[StrictInt]
+    name: Optional[StrictStr]
+    default: Any
+
+    class Config:
+        extra = Extra.ignore
+
+    def __hash__(self):
+        return hash(f"{self.id}{self.group_id}{self.name}")
+
+
+class WMAttributeGroup(TimedBaseModel):
+    id: Optional[StrictInt]
+    group_type: Optional[WMGroupTypeEnum]
+    class_id: Optional[StrictInt]
+    name: Optional[StrictStr]
+    isRequired: bool = Field(default=False, alias="is_required")
+    attributes: Optional[List[WMAttribute]]
+    default_value: Any
+
+    class Config:
+        extra = Extra.ignore
+
+    def __hash__(self):
+        return hash(f"{self.id}{self.class_id}{self.name}")
+
+
+class WMAnnotationClassEntity(TimedBaseModel):
+    id: Optional[StrictInt]
+    project_id: Optional[StrictInt]
+    type: WMClassTypeEnum = WMClassTypeEnum.OBJECT
+    name: StrictStr
+    color: HexColor
+    attribute_groups: List[WMAttributeGroup] = Field(
+        default=[], alias="attributeGroups"
+    )
+
+    def __hash__(self):
+        return hash(f"{self.id}{self.type}{self.name}")
+
+    class Config:
+        extra = Extra.ignore
+        json_encoders = {
+            HexColor: lambda v: v.__root__,
+            # WMClassTypeEnum: lambda v: v.name,
+        }
+        validate_assignment = True
+
+    @validator("type", pre=True)
+    def validate_type(cls, v):
+        if isinstance(v, WMClassTypeEnum):
+            return v
+        if isinstance(v, str):
+            # Try by value first (e.g., "object")
+            for member in WMClassTypeEnum:
+                if member.value == v:
+                    return member
+            # Try by name (e.g., "OBJECT")
+            try:
+                return WMClassTypeEnum[v.upper()]
+            except KeyError:
+                pass
+        raise ValueError(f"Invalid type: {v}")

src/superannotate/lib/core/enums.py

@@ -175,6 +175,20 @@ def get_value(cls, name):
         return cls.OBJECT.value
 
 
+class WMClassTypeEnum(Enum):
+    OBJECT = "object"
+    TAG = "tag"
+    RELATIONSHIP = "relationship"
+
+
+class WMGroupTypeEnum(Enum):
+    RADIO = "radio"
+    CHECKLIST = "checklist"
+    NUMERIC = "numeric"
+    TEXT = "text"
+    OCR = "ocr"
+
+
 class IntegrationTypeEnum(BaseTitledEnum):
     AWS = "aws", 1
     GCP = "gcp", 2
@@ -193,7 +207,7 @@ class TrainingStatus(BaseTitledEnum):
     FAILED_AFTER_EVALUATION_WITH_SAVE_MODEL = "FailedAfterEvaluationWithSavedModel", 6
 
 
-class CustomFieldEntityEnum(str, Enum):
+class CustomFieldEntityEnum(Enum):
     CONTRIBUTOR = "Contributor"
     TEAM = "Team"
     PROJECT = "Project"
@@ -207,6 +221,6 @@ class CustomFieldType(Enum):
     NUMERIC = 5
 
 
-class WMUserStateEnum(str, Enum):
+class WMUserStateEnum(Enum):
     Pending = "PENDING"
     Confirmed = "CONFIRMED"

src/superannotate/lib/core/service_types.py

@@ -138,6 +138,10 @@ def __str__(self):
         return f"Status: {self.status_code}, Error {self.error}"
 
 
+class WMClassesResponse(ServiceResponse):
+    res_data: List[entities.WMAnnotationClassEntity] = None
+
+
 class BaseItemResponse(ServiceResponse):
     res_data: entities.BaseItemEntity = None
 
@@ -186,6 +190,10 @@ class IntegrationListResponse(ServiceResponse):
     res_data: _IntegrationResponse
 
 
+class AnnotationClassResponse(ServiceResponse):
+    res_data: entities.AnnotationClassEntity = None
+
+
 class AnnotationClassListResponse(ServiceResponse):
     res_data: List[entities.AnnotationClassEntity] = None