[CI] Improved mypy typing

.circleci/config.yml

@@ -207,8 +207,8 @@ jobs:
               sudo rm  llvm-config
               sudo ln -s llvm-config-10 llvm-config
               cd ~/project/
-              pip install -U black flake8 flake8-bugbear flake8-builtins flake8-comprehensions flake8-return flake8-simplify "isort[pyproject]" numpy pytest sphinx pillow tqdm hypothesis
-              pip install -r requirements.txt --progress-bar off
+              pip install -U black flake8 flake8-bugbear flake8-builtins flake8-comprehensions flake8-return flake8-simplify "isort[pyproject]" mypy numpy pytest sphinx pillow tqdm hypothesis
+              pip install -r requirements.txt torch --progress-bar off
       - run:
           name: run black
           command: |
@@ -226,6 +226,9 @@ jobs:
           command: |
               flake8 --version
               flake8 habitat_sim/. examples/. tests/. setup.py
+      - run:
+          name: run mypy
+          command: mypy
   cpp_lint:
     docker:
       - image: circleci/buildpack-deps:18.04

habitat_sim/agent/agent.py

@@ -9,7 +9,7 @@
 import attr
 import magnum as mn
 import numpy as np
-import quaternion
+import quaternion as qt
 
 import habitat_sim.errors
 from habitat_sim import bindings as hsim
@@ -57,8 +57,8 @@ def _triple_zero() -> np.ndarray:
     return np.zeros(3)
 
 
-def _default_quaternion() -> np.quaternion:
-    return quaternion.quaternion(1, 0, 0, 0)
+def _default_quaternion() -> qt.quaternion:
+    return qt.quaternion(1, 0, 0, 0)
 
 
 @attr.s(auto_attribs=True, slots=True)
@@ -70,15 +70,15 @@ class SixDOFPose:
     """
 
     position: np.ndarray = attr.ib(factory=_triple_zero, validator=all_is_finite)
-    rotation: Union[np.quaternion, List] = attr.ib(
+    rotation: Union[qt.quaternion, List] = attr.ib(
         factory=_default_quaternion, validator=is_unit_length
     )
 
 
 @attr.s(auto_attribs=True, slots=True)
 class AgentState:
     position: np.ndarray = attr.ib(factory=_triple_zero, validator=all_is_finite)
-    rotation: Union[np.quaternion, List, np.ndarray] = attr.ib(
+    rotation: Union[qt.quaternion, List, np.ndarray] = attr.ib(
         factory=_default_quaternion, validator=is_unit_length
     )
     velocity: np.ndarray = attr.ib(factory=_triple_zero, validator=all_is_finite)

habitat_sim/simulator.py

@@ -38,6 +38,7 @@
 from habitat_sim.utils.common import quat_from_angle_axis
 
 # TODO maybe clean up types with TypeVars
+ObservationDict = Dict[str, Union[bool, np.ndarray, "Tensor"]]
 
 
 @attr.s(auto_attribs=True, slots=True)
@@ -146,23 +147,16 @@ def seed(self, new_seed: int) -> None:
         self.pathfinder.seed(new_seed)
 
     @overload
-    def reset(
-        self, agent_ids: List[int]
-    ) -> Dict[int, Dict[str, Union[ndarray, "Tensor"]]]:
+    def reset(self, agent_ids: List[int]) -> Dict[int, ObservationDict]:
         ...
 
     @overload
-    def reset(
-        self, agent_ids: Optional[int] = None
-    ) -> Dict[str, Union[ndarray, "Tensor"]]:
+    def reset(self, agent_ids: Optional[int] = None) -> ObservationDict:
         ...
 
     def reset(
         self, agent_ids: Union[Optional[int], List[int]] = None
-    ) -> Union[
-        Dict[str, Union[ndarray, "Tensor"]],
-        Dict[int, Dict[str, Union[ndarray, "Tensor"]]],
-    ]:
+    ) -> Union[ObservationDict, Dict[int, ObservationDict],]:
         super().reset()
         for i in range(len(self.agents)):
             self.reset_agent(i)
@@ -317,23 +311,18 @@ def initialize_agent(
         return agent
 
     @overload
-    def get_sensor_observations(
-        self, agent_ids: int = 0
-    ) -> Dict[str, Union[ndarray, "Tensor"]]:
+    def get_sensor_observations(self, agent_ids: int = 0) -> ObservationDict:
         ...
 
     @overload
     def get_sensor_observations(
         self, agent_ids: List[int]
-    ) -> Dict[int, Dict[str, Union[ndarray, "Tensor"]]]:
+    ) -> Dict[int, ObservationDict]:
         ...
 
     def get_sensor_observations(
         self, agent_ids: Union[int, List[int]] = 0
-    ) -> Union[
-        Dict[str, Union[ndarray, "Tensor"]],
-        Dict[int, Dict[str, Union[ndarray, "Tensor"]]],
-    ]:
+    ) -> Union[ObservationDict, Dict[int, ObservationDict],]:
         if isinstance(agent_ids, int):
             agent_ids = [agent_ids]
             return_single = True
@@ -346,9 +335,9 @@ def get_sensor_observations(
                 sensor.draw_observation()
 
         # As backport. All Dicts are ordered in Python >= 3.7
-        observations: Dict[int, Dict[str, Union[ndarray, "Tensor"]]] = OrderedDict()
+        observations: Dict[int, ObservationDict] = OrderedDict()
         for agent_id in agent_ids:
-            agent_observations: Dict[str, Union[ndarray, "Tensor"]] = {}
+            agent_observations: ObservationDict = {}
             for sensor_uuid, sensor in self.__sensors[agent_id].items():
                 agent_observations[sensor_uuid] = sensor.get_observation()
             observations[agent_id] = agent_observations
@@ -382,25 +371,20 @@ def last_state(self, agent_id: Optional[int] = None) -> AgentState:
         return self.__last_state[agent_id]
 
     @overload
-    def step(
-        self, action: Union[str, int], dt: float = 1.0 / 60.0
-    ) -> Dict[str, Union[bool, ndarray, "Tensor"]]:
+    def step(self, action: Union[str, int], dt: float = 1.0 / 60.0) -> ObservationDict:
         ...
 
     @overload
     def step(
         self, action: MutableMapping_T[int, Union[str, int]], dt: float = 1.0 / 60.0
-    ) -> Dict[int, Dict[str, Union[bool, ndarray, "Tensor"]]]:
+    ) -> Dict[int, ObservationDict]:
         ...
 
     def step(
         self,
         action: Union[str, int, MutableMapping_T[int, Union[str, int]]],
         dt: float = 1.0 / 60.0,
-    ) -> Union[
-        Dict[str, Union[bool, ndarray, "Tensor"]],
-        Dict[int, Dict[str, Union[bool, ndarray, "Tensor"]]],
-    ]:
+    ) -> Union[ObservationDict, Dict[int, ObservationDict],]:
         self._num_total_frames += 1
         if isinstance(action, MutableMapping):
             return_single = False
@@ -478,6 +462,7 @@ class Sensor:
 
     TODO(MS) define entire Sensor class in python, reducing complexity
     """
+    buffer = Union[np.ndarray, "Tensor"]
 
     def __init__(self, sim: Simulator, agent: Agent, sensor_id: str) -> None:
         self._sim = sim
@@ -499,7 +484,7 @@ def __init__(self, sim: Simulator, agent: Agent, sensor_id: str) -> None:
 
             resolution = self._spec.resolution
             if self._spec.sensor_type == SensorType.SEMANTIC:
-                self._buffer = torch.empty(
+                self._buffer: Union[np.ndarray, "Tensor"] = torch.empty(
                     resolution[0], resolution[1], dtype=torch.int32, device=device
                 )
             elif self._spec.sensor_type == SensorType.DEPTH:
@@ -560,15 +545,15 @@ def get_observation(self) -> Union[ndarray, "Tensor"]:
         tgt = self._sensor_object.render_target
 
         if self._spec.gpu2gpu_transfer:
-            with torch.cuda.device(self._buffer.device):  # type: ignore[attr-defined]
+            with torch.cuda.device(self._buffer.device):  # type: ignore[attr-defined, union-attr]
                 if self._spec.sensor_type == SensorType.SEMANTIC:
-                    tgt.read_frame_object_id_gpu(self._buffer.data_ptr())  # type: ignore[attr-defined]
+                    tgt.read_frame_object_id_gpu(self._buffer.data_ptr())  # type: ignore[attr-defined, union-attr]
                 elif self._spec.sensor_type == SensorType.DEPTH:
-                    tgt.read_frame_depth_gpu(self._buffer.data_ptr())  # type: ignore[attr-defined]
+                    tgt.read_frame_depth_gpu(self._buffer.data_ptr())  # type: ignore[attr-defined, union-attr]
                 else:
-                    tgt.read_frame_rgba_gpu(self._buffer.data_ptr())  # type: ignore[attr-defined]
+                    tgt.read_frame_rgba_gpu(self._buffer.data_ptr())  # type: ignore[attr-defined, union-attr]
 
-                obs = self._buffer.flip(0)
+                obs = self._buffer.flip(0)  # type: ignore[union-attr]
         else:
             size = self._sensor_object.framebuffer_size
 

habitat_sim/utils/common.py

@@ -6,29 +6,29 @@
 
 import math
 from io import BytesIO
-from typing import List, Sequence, Tuple
+from typing import List, Sequence, Tuple, Union
 from urllib.request import urlopen
 from zipfile import ZipFile
 
 import magnum as mn
 import numpy as np
-import quaternion
+import quaternion as qt
 
 
-def quat_from_coeffs(coeffs: Sequence[float]) -> np.quaternion:
+def quat_from_coeffs(coeffs: Union[Sequence[float], np.ndarray]) -> qt.quaternion:
     r"""Creates a quaternion from the coeffs returned by the simulator backend
 
     :param coeffs: Coefficients of a quaternion in :py:`[b, c, d, a]` format,
         where :math:`q = a + bi + cj + dk`
     :return: A quaternion from the coeffs
     """
-    quat = np.quaternion(1, 0, 0, 0)
+    quat = qt.quaternion(1, 0, 0, 0)
     quat.real = coeffs[3]
     quat.imag = coeffs[0:3]
     return quat
 
 
-def quat_to_coeffs(quat: np.quaternion) -> np.ndarray:
+def quat_to_coeffs(quat: qt.quaternion) -> np.ndarray:
     r"""Converts a quaternion into the coeffs format the backend expects
 
     :param quat: The quaternion
@@ -41,18 +41,18 @@ def quat_to_coeffs(quat: np.quaternion) -> np.ndarray:
     return coeffs
 
 
-def quat_to_magnum(quat: np.quaternion) -> mn.Quaternion:
+def quat_to_magnum(quat: qt.quaternion) -> mn.Quaternion:
     return mn.Quaternion(quat.imag, quat.real)
 
 
-def quat_from_magnum(quat: mn.Quaternion) -> np.quaternion:
-    a = np.quaternion(1, 0, 0, 0)
+def quat_from_magnum(quat: mn.Quaternion) -> qt.quaternion:
+    a = qt.quaternion(1, 0, 0, 0)
     a.real = quat.scalar
     a.imag = quat.vector
     return a
 
 
-def quat_to_angle_axis(quat: np.quaternion) -> Tuple[float, np.ndarray]:
+def quat_to_angle_axis(quat: qt.quaternion) -> Tuple[float, np.ndarray]:
     r"""Converts a quaternion to angle axis format
 
     :param quat: The quaternion
@@ -62,7 +62,7 @@ def quat_to_angle_axis(quat: np.quaternion) -> Tuple[float, np.ndarray]:
             then this is harded coded to be the +x axis
     """
 
-    rot_vec = quaternion.as_rotation_vector(quat)
+    rot_vec = qt.as_rotation_vector(quat)
 
     theta = np.linalg.norm(rot_vec)
     if np.abs(theta) < 1e-5:
@@ -74,19 +74,19 @@ def quat_to_angle_axis(quat: np.quaternion) -> Tuple[float, np.ndarray]:
     return (theta, w)
 
 
-def quat_from_angle_axis(theta: float, axis: np.ndarray) -> np.quaternion:
+def quat_from_angle_axis(theta: float, axis: np.ndarray) -> qt.quaternion:
     r"""Creates a quaternion from angle axis format
 
     :param theta: The angle to rotate about the axis by
     :param axis: The axis to rotate about
     :return: The quaternion
     """
-    axis = axis.astype(np.float)
+    axis = axis.astype(float)
     axis /= np.linalg.norm(axis)
-    return quaternion.from_rotation_vector(theta * axis)
+    return qt.from_rotation_vector(theta * axis)
 
 
-def quat_from_two_vectors(v0: np.ndarray, v1: np.ndarray) -> np.quaternion:
+def quat_from_two_vectors(v0: np.ndarray, v1: np.ndarray) -> qt.quaternion:
     r"""Creates a quaternion that rotates the first vector onto the second vector
 
     :param v0: The starting vector, does not need to be a unit vector
@@ -111,14 +111,14 @@ def quat_from_two_vectors(v0: np.ndarray, v1: np.ndarray) -> np.quaternion:
         w2 = (1 + c) * 0.5
         w = np.sqrt(w2)
         axis = axis * np.sqrt(1 - w2)
-        return np.quaternion(w, *axis)
+        return qt.quaternion(w, *axis)
 
     axis = np.cross(v0, v1)
     s = np.sqrt((1 + c) * 2)
-    return np.quaternion(s * 0.5, *(axis / s))
+    return qt.quaternion(s * 0.5, *(axis / s))
 
 
-def angle_between_quats(q1: np.quaternion, q2: np.quaternion) -> float:
+def angle_between_quats(q1: qt.quaternion, q2: qt.quaternion) -> float:
     r"""Computes the angular distance between two quaternions
 
     :return: The angular distance between q1 and q2 in radians
@@ -130,7 +130,7 @@ def angle_between_quats(q1: np.quaternion, q2: np.quaternion) -> float:
     return 2 * np.arctan2(np.linalg.norm(dq.imag), np.abs(dq.real))
 
 
-def quat_rotate_vector(q: np.quaternion, v: np.ndarray) -> np.ndarray:
+def quat_rotate_vector(q: qt.quaternion, v: np.ndarray) -> np.ndarray:
     r"""Helper function to rotate a vector by a quaternion
 
     :param q: The quaternion to rotate the vector with
@@ -141,10 +141,10 @@ def quat_rotate_vector(q: np.quaternion, v: np.ndarray) -> np.ndarray:
 
     .. code:: py
 
-        v = (q * np.quaternion(0, *v) * q.inverse()).imag
+        v = (q * qt.quaternion(0, *v) * q.inverse()).imag
     """
 
-    vq = np.quaternion(0, 0, 0, 0)
+    vq = qt.quaternion(0, 0, 0, 0)
     vq.imag = v
     return (q * vq * q.inverse()).imag
 

habitat_sim/utils/data/pose_extractor.py

@@ -89,7 +89,7 @@ def _convert_to_scene_coordinate_system(
         self,
         poses: List[Tuple[Tuple[int, int], Tuple[int, int], str]],
         ref_point: Tuple[float32, float32, float32],
-    ) -> List[Tuple[Union[np.ndarray, Tuple[int, int]], quaternion, str]]:
+    ) -> List[Tuple[Tuple[int, int], quaternion, str]]:
         # Convert from topdown map coordinate system to that of the scene
         startw, starty, starth = ref_point
         for i, pose in enumerate(poses):
@@ -112,7 +112,8 @@ def _convert_to_scene_coordinate_system(
             )
             cam_normal = new_cpi - new_pos
             new_rot = self._compute_quat(cam_normal)
-            poses[i] = (new_pos, new_rot, filepath)
+            new_pos_t: Tuple[int, int] = tuple(new_pos)  # type: ignore[assignment]
+            poses[i] = (new_pos_t, new_rot, filepath)
 
         return poses
 

habitat_sim/utils/viz_utils.py

@@ -266,7 +266,7 @@ def depth_to_rgb(depth_image: np.ndarray, clip_max: float = 10.0) -> np.ndarray:
     d_im = np.clip(depth_image, 0, clip_max)
     d_im /= clip_max
     rgb_d_im = (d_im * 255).astype(np.uint8)
-    return rgb_d_im
+    return np.asarray(rgb_d_im)
 
 
 def semantic_to_rgb(semantic_image: np.ndarray) -> np.ndarray:

tests/test_controls.py

@@ -9,7 +9,7 @@
 import magnum as mn
 import numpy as np
 import pytest
-import quaternion  # noqa: F401
+import quaternion as qt
 from hypothesis import strategies as st
 
 import habitat_sim
@@ -48,7 +48,7 @@ def test_no_move_fun():
 @attr.s(auto_attribs=True, cmp=False)
 class ExpectedDelta:
     delta_pos: np.ndarray = attr.Factory(lambda: np.array([0, 0, 0]))
-    delta_rot: np.quaternion = attr.Factory(lambda: np.quaternion(1, 0, 0, 0))
+    delta_rot: qt.quaternion = attr.Factory(lambda: qt.quaternion(1, 0, 0, 0))
 
 
 def _check_state_same(s1, s2):

tests/test_pyrobot_noisy_controls.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 import pytest
-import quaternion  # noqa: F401
+import quaternion as qt
 
 import habitat_sim
 import habitat_sim.errors
@@ -89,7 +89,7 @@ def test_pyrobot_noisy_actions(noise_multiplier, robot, controller):
 
     for base_action in {act.replace("noisy_", "") for act in agent_config.action_space}:
         state = agent.state
-        state.rotation = np.quaternion(1, 0, 0, 0)
+        state.rotation = qt.quaternion(1, 0, 0, 0)
         agent.state = state
         agent.act(base_action)
         base_state = agent.state