Merge pull request #836 from superannotateai/fix_methods

fix related pydantic

Author: VaghinakDev

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

@@ -31,15 +31,21 @@ def _validate_hex_color(v: str) -> str:
 HexColor = Annotated[str, AfterValidator(_validate_hex_color)]
 
 
-def _validate_string_date(v: datetime) -> str:
-    """Convert datetime to string format."""
+def _validate_string_date(v: Union[datetime, str]) -> str:
+    """Convert datetime to string format. For case data output."""
     if isinstance(v, str):
-        return v
-    return v.isoformat().split("+")[0] + ".000Z"
+        try:
+            dt = datetime.fromisoformat(v.replace("Z", "+00:00"))
+            return dt.strftime("%Y-%m-%dT%H:%M:%S+00:00")
+        except (ValueError, AttributeError):
+            return v
+    elif isinstance(v, datetime):
+        return v.strftime("%Y-%m-%dT%H:%M:%S+00:00")
+    return v
 
 
-def _serialize_string_date(v) -> str:
-    """Serialize datetime or string to string format."""
+def _serialize_string_date(v: Union[datetime, str]) -> str:
+    """Serialize datetime or string to string format. For case data input."""
     if isinstance(v, str):
         return v
     if isinstance(v, datetime):
@@ -127,7 +133,7 @@ def _validate_token(value: str) -> str:
 class ConfigEntity(BaseModel):
     model_config = ConfigDict(extra="ignore")
 
-    API_TOKEN: str = Field(alias="SA_TOKEN")
+    API_TOKEN: TokenStr = Field(alias="SA_TOKEN")
     API_URL: str = Field(alias="SA_URL", default=BACKEND_URL)
     LOGGING_LEVEL: Literal[
         "NOTSET", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"

src/superannotate/lib/core/enums.py

@@ -2,6 +2,10 @@
 from enum import Enum
 from types import DynamicClassAttribute
 
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+from pydantic_core import PydanticCustomError
+
 
 class classproperty:  # noqa
     def __init__(self, getter):
@@ -33,6 +37,47 @@ def value(self) -> int:
     def __unicode__(self):
         return self.__doc__
 
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type, handler: GetCoreSchemaHandler):
+        """Customize Pydantic v2 validation to show titles in error messages."""
+        return core_schema.no_info_after_validator_function(
+            cls._validate,
+            core_schema.union_schema(
+                [
+                    core_schema.is_instance_schema(cls),
+                    core_schema.int_schema(),
+                    core_schema.str_schema(),
+                ]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda x: x.value if isinstance(x, cls) else x, when_used="json"
+            ),
+        )
+
+    @classmethod
+    def _validate(cls, value):
+        """Validate and convert value to enum member."""
+        if isinstance(value, cls):
+            return value
+
+        # Try to find by value
+        if isinstance(value, int):
+            for enum in cls:
+                if enum.value == value:
+                    return enum
+
+        # Try to find by title or name
+        if isinstance(value, str):
+            for enum in cls:
+                if enum.__doc__ and enum.__doc__.lower() == value.lower():
+                    return enum
+            if value in cls.__members__:
+                return cls.__members__[value]
+
+        # Build error message with titles
+        available = ", ".join(f"{enum.__doc__.lower()}" for enum in cls if enum.__doc__)
+        raise PydanticCustomError("enum", f"Input should be: {available}")
+
     @classmethod
     def choices(cls) -> typing.Tuple[str]:
         """Return all titles as choices."""

tests/integration/classes/test_create_update_annotation_class.py

@@ -658,6 +658,14 @@ def test_update_annotation_class_without_group_type(self):
             )
             assert res["attribute_groups"][0]["group_type"] == "radio"
 
+    def test_create_with_invalid_type(self):
+        try:
+            sa.create_annotation_class(
+                self.PROJECT_NAME, "tt", "#FFFFFF", class_type="invalid"
+            )
+        except AppException as e:
+            assert "Input should be: object, tag, relationship" in str(e)
+
 
 class TestVideoCreateAnnotationClasses(BaseTestCase):
     PROJECT_NAME = "TestVideoCreateAnnotationClasses"

tests/unit/test_init.py

@@ -16,7 +16,7 @@ class ClientInitTestCase(TestCase):
 
     def test_init_via_invalid_token(self):
         _token = "123"
-        with self.assertRaisesRegex(AppException, r"(\s+)token(\s+)Invalid token."):
+        with self.assertRaisesRegex(AppException, r"Invalid token\."):
             SAClient(token=_token)
 
     @patch("lib.infrastructure.controller.Controller.get_current_user")
@@ -61,9 +61,7 @@ def test_init_via_config_json_invalid_json(self):
                 with open(f"{config_dir}/config.json", "w") as config_json:
                     json.dump({"token": "INVALID_TOKEN"}, config_json)
                 for kwargs in ({}, {"config_path": f"{config_dir}/config.json"}):
-                    with self.assertRaisesRegex(
-                        AppException, r"(\s+)token(\s+)Invalid token."
-                    ):
+                    with self.assertRaisesRegex(AppException, r"Invalid token\."):
                         SAClient(**kwargs)
 
     @patch("lib.infrastructure.controller.Controller.get_current_user")
@@ -137,7 +135,7 @@ def test_init_env(self, get_team, get_current_user):
 
     @patch.dict(os.environ, {"SA_URL": "SOME_URL", "SA_TOKEN": "SOME_TOKEN"})
     def test_init_env_invalid_token(self):
-        with self.assertRaisesRegex(AppException, r"(\s+)SA_TOKEN(\s+)Invalid token."):
+        with self.assertRaisesRegex(AppException, r"Invalid token\."):
             SAClient()
 
     def test_init_via_config_ini_invalid_token(self):
@@ -157,9 +155,7 @@ def test_init_via_config_ini_invalid_token(self):
                     config_parser.write(config_ini)
 
                 for kwargs in ({}, {"config_path": f"{config_dir}/config.ini"}):
-                    with self.assertRaisesRegex(
-                        AppException, r"(\s+)SA_TOKEN(\s+)Invalid token."
-                    ):
+                    with self.assertRaisesRegex(AppException, r"Invalid token\."):
                         SAClient(**kwargs)
 
     def test_invalid_config_path(self):