Fix mypy in a more neat way

Signed-off-by: Kim, Vinnam <[email protected]>
openvinotoolkit · Apr 7, 2023 · c5d476b · c5d476b
1 parent 774bf99
commit c5d476b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/otx/api/entities/dataset_item.py b/otx/api/entities/dataset_item.py
@@ -10,7 +10,7 @@
 import itertools
 import logging
 from threading import Lock
-from typing import List, Optional, Sequence, Set, Tuple, Union
+from typing import List, Optional, Sequence, Set, Tuple, TypeVar, Union
 from bson import ObjectId
 import numpy as np
 
@@ -28,6 +28,9 @@
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T", bound="DatasetItemEntity")
+
+
 class DatasetItemEntity(metaclass=abc.ABCMeta):
     """DatasetItemEntity represents an item in the DatasetEntity.
 
@@ -488,9 +491,9 @@ def get_metadata_by_name_and_model(self, name: str, model: Optional[ModelEntity]
         """
         return [meta for meta in self.get_metadata() if meta.data.name == name and meta.model == model]
 
-    def wrap(self, **kwargs) -> "DatasetItemEntity":
+    def wrap(self: T, **kwargs) -> T:
         """Creates a new DatasetItemEntity, overriding only the given arguments to the existing ones for this instance."""
-        return DatasetItemEntity(
+        return self.__class__(
             media=kwargs.get("media", self.media),
             annotation_scene=kwargs.get("annotation_scene", self.annotation_scene),
             roi=kwargs.get("roi", self.roi),
@@ -538,7 +541,7 @@ def __eq__(self, other):
 
     def wrap(self, **kwargs) -> "DatasetItemEntityWithID":
         """Creates a new DatasetItemEntityWithID, overriding only the given arguments to the existing ones for this instance."""
-        return DatasetItemEntityWithID(
+        return self.__class__(
             media=kwargs.get("media", self.media),
             annotation_scene=kwargs.get("annotation_scene", self.annotation_scene),
             roi=kwargs.get("roi", self.roi),

diff --git a/otx/api/entities/datasets.py b/otx/api/entities/datasets.py
@@ -82,7 +82,7 @@ def __next__(self) -> DatasetItemEntity:
         return item
 
 
-TDatasetItemEntity = TypeVar("TDatasetItemEntity", bound=DatasetItemEntity)
+TDatasetItemEntity = TypeVar("TDatasetItemEntity", bound="DatasetItemEntity")
 
 
 class DatasetEntity(Generic[TDatasetItemEntity]):
@@ -309,9 +309,10 @@ def with_empty_annotations(
                 Defaults to AnnotationSceneKind.PREDICTION
 
         Returns:
-            DatasetEntity: a new dataset containing the same items, with empty annotation objects.
+            DatasetEntity: a new dataset containing the same items, with empty annotation objects.asdf
         """
-        new_dataset = DatasetEntity[DatasetItemEntity](purpose=self.purpose)
+        new_dataset = DatasetEntity[TDatasetItemEntity](purpose=self.purpose)
+
         for dataset_item in self:
             if isinstance(dataset_item, DatasetItemEntity):
                 empty_annotation = AnnotationSceneEntity(annotations=[], kind=annotation_kind)
@@ -360,7 +361,7 @@ def remove(self, item: TDatasetItemEntity) -> None:
         This function calls remove_at_indices function.
 
         Args:
-            item (TDatasetItemEntity): the item to be deleted.
+            item (DatasetItemEntity): the item to be deleted.
 
         Raises:
             ValueError: if the input item is not in the dataset
@@ -384,7 +385,7 @@ def append(self, item: TDatasetItemEntity) -> None:
             >>> dataset.append(dataset_item)
 
         Args:
-            item (TDatasetItemEntity): item to append
+            item (DatasetItemEntity): item to append
         """
 
         if item.media is None: