Changed return type of create_entity and create_multiple_entities to robofish.io.Entity

src/robofish/io/entity.py

@@ -30,7 +30,7 @@ class Entity(h5py.Group):
         orientations: Iterable = None,
         outlines: Iterable = None,
         sampling: str = None,
-    ):
+    ) -> robofish.io.Entity:
         poses, positions, orientations, outlines = utils.np_array(
             poses, positions, orientations, outlines
         )

src/robofish/io/file.py

@@ -439,7 +439,7 @@ class File(h5py.File):
         orientations: Iterable = None,
         outlines: Iterable = None,
         sampling: str = None,
-    ) -> str:
+    ) -> robofish.io.Entity:
         """Creates a new single entity.
 
         Args:
@@ -478,7 +478,7 @@ class File(h5py.File):
         names: Iterable[str] = None,
         outlines=None,
         sampling=None,
-    ) -> Iterable:
+    ) -> Iterable[robofish.io.Entity]:
         """Creates multiple entities.
 
         Args:
@@ -496,7 +496,7 @@ class File(h5py.File):
         ), f"A 3 dimensional array was expected (entity, timestep, 3). There were {poses.ndim} dimensions in poses: {poses.shape}"
         assert poses.shape[2] in [3, 4]
         agents = poses.shape[0]
-        entity_names = []
+        entities = []
 
         for i in range(agents):
             e_name = None if names is None else names[i]
@@ -504,7 +504,7 @@ class File(h5py.File):
                 outlines if outlines is None or outlines.ndim == 3 else outlines[i]
             )
 
-            entity_names.append(
+            entities.append(
                 self.create_entity(
                     category=category,
                     sampling=sampling,
@@ -513,7 +513,7 @@ class File(h5py.File):
                     outlines=e_outline,
                 )
             )
-        return entity_names
+        return entities
 
     def update_calculated_data(self, verbose=False):
         changed = any([e.update_calculated_data(verbose) for e in self.entities])