Merge pull request #9 from rolisz/master

RequestModelFactory should look at Config.title attribute
doccano · May 18, 2021 · 3203f6d · 3203f6d
2 parents 32d1785 + ff5ec02
commit 3203f6d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/auto_labeling_pipeline/models.py b/auto_labeling_pipeline/models.py
@@ -25,7 +25,7 @@ def create(cls, model_name: str, attributes: Dict) -> RequestModel:
     @classmethod
     def find(cls, model_name: str) -> Type[RequestModel]:
         for subclass in RequestModel.__subclasses__():
-            if subclass.__name__ == model_name:
+            if subclass.Config.title == model_name:
                 return subclass
         raise NameError(f'{model_name} is not found.')
 

diff --git a/tests/test_models.py b/tests/test_models.py
@@ -18,7 +18,7 @@ def test_request_model_factory_creates_model_correctly():
         type='PLAIN_TEXT',
         language='en'
     )
-    model_name = model.__repr_name__()
+    model_name = model.Config.title
     attributes = model.dict()
     restored_model = RequestModelFactory.create(model_name, attributes)
     assert restored_model == model