Enums

Author: VaghinakDev

pytest.ini

@@ -2,4 +2,4 @@
 minversion = 3.7
 log_cli=true
 python_files = test_*.py
-addopts = -n auto --dist=loadscope
+;addopts = -n auto --dist=loadscope

src/superannotate/__init__.py

@@ -18,6 +18,7 @@
 from superannotate.lib.core import PACKAGE_VERSION_UPGRADE  # noqa
 from superannotate.logger import get_default_logger  # noqa
 from superannotate.version import __version__  # noqa
+import superannotate.lib.core.enums as enums
 
 
 SESSIONS = {}
@@ -26,6 +27,7 @@
     "__version__",
     "SAClient",
     # Utils
+    "enums",
     "AppException",
     # analytics
     "class_distribution",

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

@@ -22,10 +22,23 @@
 from pydantic import validate_arguments as pydantic_validate_arguments
 from pydantic import ValidationError
 from pydantic.errors import StrRegexError
+from pydantic import errors
+from pydantic.errors import PydanticTypeError
 
 NotEmptyStr = constr(strict=True, min_length=1)
 
 
+class EnumMemberError(PydanticTypeError):
+    code = 'enum'
+
+    def __str__(self) -> str:
+        permitted = ', '.join(str(v.name) for v in self.enum_values)  # type: ignore
+        return f'Available values are: {permitted}'
+
+
+errors.EnumMemberError = EnumMemberError
+
+
 class EmailStr(StrictStr):
     @classmethod
     def validate(cls, value: Union[str]) -> Union[str]:

src/superannotate/lib/core/enums.py

@@ -8,8 +8,13 @@ def __new__(cls, title, value):
         obj._value_ = value
         obj.__doc__ = title
         obj._type = "titled_enum"
+        cls._value2member_map_[title] = obj
         return obj
 
+    @classmethod
+    def choices(cls):
+        return tuple(cls._value2member_map_.keys())
+
     @DynamicClassAttribute
     def name(self) -> str:
         return self.__doc__
@@ -41,6 +46,9 @@ def titles(cls):
     def equals(self, other: Enum):
         return self.__doc__.lower() == other.__doc__.lower()
 
+    def __eq__(self, other):
+        return super(BaseTitledEnum, self).__eq__(other)
+
 
 class AnnotationTypes(str, Enum):
     BBOX = "bbox"

tests/unit/test_enum_arguments_handeling.py

@@ -0,0 +1,20 @@
+# from typing import Literal
+from pydantic.typing import Literal
+
+from superannotate import enums
+from superannotate import SAClient
+from superannotate.lib.app.interface.types import validate_arguments
+
+
+
+@validate_arguments
+def foo(status: enums.ProjectStatus):
+    return status
+
+
+def test_enum_arg():
+    SAClient()
+    assert foo(1) == 1
+    assert foo("NotStarted") == 1
+    assert foo(enums.ProjectStatus.NotStarted.name) == 1
+    assert foo(enums.ProjectStatus.NotStarted.value) == 1