Fix some type annotations

src/robofish/io/entity.py

@@ -2,13 +2,11 @@
 .. include:: ../../../docs/entity.md
 """
 
-import robofish.io
 import robofish.io.utils as utils
 
 import h5py
 import numpy as np
-from typing import Iterable, Union
-import datetime
+from typing import Iterable, Optional
 import logging
 import deprecation
 
@@ -24,14 +22,14 @@ class Entity(h5py.Group):
         cla,
         entities_group,
         category: str,
-        poses: Iterable = None,
-        name: str = None,
-        individual_id: int = None,
-        positions: Iterable = None,
-        orientations: Iterable = None,
-        outlines: Iterable = None,
-        sampling: str = None,
-    ):
+        poses: Optional[Iterable] = None,
+        name: Optional[str] = None,
+        individual_id: Optional[int] = None,
+        positions: Optional[Iterable] = None,
+        orientations: Optional[Iterable] = None,
+        outlines: Optional[Iterable] = None,
+        sampling: Optional[str] = None,
+    ) -> "Entity":
         poses, positions, orientations, outlines = utils.np_array(
             poses, positions, orientations, outlines
         )
@@ -92,29 +90,29 @@ class Entity(h5py.Group):
         return ori_vec
 
     @property
-    def group_name(self):
+    def group_name(self) -> str:
         return super().name
 
     @property
-    def name(self):
+    def name(self) -> str:
         return self.group_name.split("/")[-1]
 
     @property
-    def category(self):
+    def category(self) -> str:
         return self.attrs["category"]
 
-    def create_outlines(self, outlines: Iterable, sampling=None):
+    def create_outlines(self, outlines: Iterable, sampling=None) -> None:
         outlines = self.create_dataset("outlines", data=outlines, dtype=np.float32)
         if sampling is not None:
             outlines.attrs["sampling"] = sampling
 
     def create_poses(
         self,
-        poses: Iterable = None,
-        positions: Iterable = None,
-        orientations: Iterable = None,
-        sampling: str = None,
-    ):
+        poses: Optional[Iterable] = None,
+        positions: Optional[Iterable] = None,
+        orientations: Optional[Iterable] = None,
+        sampling: Optional[str] = None,
+    ) -> None:
         poses, positions, orientations = utils.np_array(poses, positions, orientations)
 
         # Either poses or positions not both
@@ -147,11 +145,11 @@ class Entity(h5py.Group):
                 orientations.attrs["sampling"] = sampling
 
     @property
-    def positions(self):
+    def positions(self) -> np.ndarray:
         return self["positions"]
 
     @property
-    def orientations(self):
+    def orientations(self) -> np.ndarray:
         if not "orientations" in self:
             # If no orientation is given, the default direction is to the right
             return np.tile([0, 1], (self.positions.shape[0], 1))
@@ -205,7 +203,7 @@ class Entity(h5py.Group):
         return poses_with_calculated_orientation
 
     @property
-    def poses_hash(self):
+    def poses_hash(self) -> int:
         # The hash of h5py datasets changes each time the file is reopened.
         # Also the hash of casting the array to bytes and calculating the hash changes.
         def npsumhash(a):
@@ -220,11 +218,11 @@ class Entity(h5py.Group):
         return h
 
     @property
-    def poses(self):
+    def poses(self) -> np.ndarray:
         return np.concatenate([self.positions, self.orientations], axis=1)
 
     @property
-    def poses_rad(self):
+    def poses_rad(self) -> np.ndarray:
         return np.concatenate([self.positions, self.orientations_rad], axis=1)
 
     @property
@@ -260,7 +258,7 @@ class Entity(h5py.Group):
         turn = utils.limit_angle_range(diff[:, 2], _range=(-np.pi, np.pi))
         return np.stack([speed, turn], axis=-1)
 
-    def update_calculated_data(self, verbose=False, force_update=False):
+    def update_calculated_data(self, verbose=False, force_update=False) -> bool:
         changed = False
         if (
             "poses_hash" not in self.attrs
@@ -341,7 +339,7 @@ class Entity(h5py.Group):
         return np.stack([speed, turn], axis=-1)
 
     @property
-    def actions_speeds_turns(self):
+    def actions_speeds_turns(self) -> np.ndarray:
         if "calculated_actions_speeds_turns" in self:
             assert (
                 self.attrs["poses_hash"] == self.poses_hash
@@ -351,7 +349,7 @@ class Entity(h5py.Group):
             return self.calculate_actions_speeds_turns()
 
     @property
-    def orientations_rad(self):
+    def orientations_rad(self) -> np.ndarray:
         if "calculated_orientations_rad" in self:
             assert (
                 self.attrs["poses_hash"] == self.poses_hash

src/robofish/io/file.py

@@ -63,20 +63,20 @@ class File(h5py.File):
 
     def __init__(
         self,
-        path: Union[str, Path] = None,
+        path: Optinal[Union[str, Path]] = None,
         mode: str = "r",
         *,  # PEP 3102
-        world_size_cm: List[int] = None,
+        world_size_cm: Optional[List[int]] = None,
         world_shape: str = "rectangle",
         validate: bool = False,
         validate_when_saving: bool = True,
         strict_validate: bool = False,
         format_version: List[int] = default_format_version,
         format_url: str = default_format_url,
-        sampling_name: str = None,
-        frequency_hz: int = None,
-        monotonic_time_points_us: Iterable = None,
-        calendar_time_points: Iterable = None,
+        sampling_name: Optional[str] = None,
+        frequency_hz: Optional[int] = None,
+        monotonic_time_points_us: Optional[Iterable] = None,
+        calendar_time_points: Optional[Iterable] = None,
         open_copy: bool = False,
         validate_poses_hash: bool = True,
         calculate_data_on_close: bool = True,
@@ -332,11 +332,11 @@ class File(h5py.File):
 
     def create_sampling(
         self,
-        name: str = None,
-        frequency_hz: int = None,
-        monotonic_time_points_us: Iterable = None,
-        calendar_time_points: Iterable = None,
-        default: bool = False,
+        name: Optional[str] = None,
+        frequency_hz: Optional[int] = None,
+        monotonic_time_points_us: Optional[Iterable] = None,
+        calendar_time_points: Optional[Iterable] = None,
+        default: Optional[bool] = False,
     ):
         # Find Name for sampling if none is given
         if name is None:
@@ -465,24 +465,24 @@ class File(h5py.File):
     def create_entity(
         self,
         category: str,
-        poses: Iterable = None,
-        name: str = None,
-        individual_id: int = None,
-        positions: Iterable = None,
-        orientations: Iterable = None,
-        outlines: Iterable = None,
-        sampling: str = None,
-    ) -> str:
+        poses: Optional[Iterable] = None,
+        name: Optional[str] = None,
+        individual_id: Optional[int] = None,
+        positions: Optional[Iterable] = None,
+        orientations: Optional[Iterable] = None,
+        outlines: Optional[Iterable] = None,
+        sampling: Optional[str] = None,
+    ) -> Entity:
         """Creates a new single entity.
 
         Args:
             TODO
             category: the  of the entity. The canonical values are ['organism', 'robot', 'obstacle'].
             poses: optional two dimensional array, containing the poses of the entity (x,y,orientation_x, orientation_y).
-            poses_rad: optional two dimensional containing the poses of the entity (x,y, orientation_rad).
             name: optional name of the entity. If no name is given, the  is used with an id (e.g. 'fish_1')
             individual_id (int, optional): invididual id of the entity.
             outlines: optional three dimensional array, containing the outlines of the entity
+            sampling: The string refference to the sampling. If none is given, the standard sampling from creating the file is used.
         Returns:
             Name of the created entity
         """
@@ -510,11 +510,11 @@ class File(h5py.File):
         self,
         category: str,
         poses: Iterable,
-        names: Iterable[str] = None,
-        individual_ids: Iterable[int] = None,
+        names: Optional[Iterable[str]] = None,
+        individual_ids: Optional[Iterable[int]] = None,
         outlines=None,
         sampling=None,
-    ) -> Iterable:
+    ) -> Iterable[Entity]:
         """Creates multiple entities.
 
         Args:
@@ -868,13 +868,13 @@ class File(h5py.File):
 
     def plot(
         self,
-        ax: matplotlib.axes = None,
+        ax: Optional[matplotlib.axes] = None,
         lw_distances: bool = False,
         lw: int = 2,
         ms: int = 32,
-        figsize: Tuple[int] = None,
+        figsize: Optional[Tuple[int]] = None,
         step_size: int = 25,
-        c: List = None,
+        c: Optional[List] = None,
         cmap: matplotlib.colors.Colormap = "Set1",
         skip_timesteps=0,
         max_timesteps=None,
@@ -985,7 +985,7 @@ class File(h5py.File):
 
         return ax
 
-    def render(self, video_path: Union[str, Path] = None, **kwargs: Dict) -> None:
+    def render(self, video_path: Optional[Union[str, Path]] = None, **kwargs: Dict) -> None:
         """Render a video of the file.
 
         The tracks are rendered in a video using matplotlib.animation.FuncAnimation.