Add mypy and static type checking to Habitat Lab

.pre-commit-config.yaml

@@ -37,7 +37,7 @@ repos:
     rev: master
     hooks:
     -   id: autoflake
-        args: ['--expand-star-imports', '--ignore-init-module-imports', '--in-place', '-c']
+        args: ['--expand-star-imports', '--ignore-init-module-imports', '--in-place']
         exclude: docs/
 
 -   repo: https://gitlab.com/pycqa/flake8
@@ -49,6 +49,12 @@ repos:
         - flake8-bugbear
         - flake8-comprehensions
 
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.782'
+    hooks:
+    -   id: mypy
+        pass_filenames: false
+
 -   repo: https://github.com/kynan/nbstripout
     rev: master
     hooks:

habitat/core/agent.py

@@ -8,9 +8,10 @@
 and ``reset()`` methods.
 """
 
-from typing import Any, Dict, Union
+from typing import TYPE_CHECKING, Any, Dict, Union
 
-from habitat.core.simulator import Observations
+if TYPE_CHECKING:
+    from habitat.core.simulator import Observations
 
 
 class Agent:
@@ -24,7 +25,7 @@ def reset(self) -> None:
         raise NotImplementedError
 
     def act(
-        self, observations: Observations
+        self, observations: "Observations"
     ) -> Union[int, str, Dict[str, Any]]:
         r"""Called to produce an action to perform in an environment.
 

habitat/core/benchmark.py

@@ -11,18 +11,20 @@
 
 import os
 from collections import defaultdict
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
 from habitat.config.default import get_config
-from habitat.core.agent import Agent
 from habitat.core.env import Env
 
+if TYPE_CHECKING:
+    from habitat.core.agent import Agent
+
 
 class Benchmark:
     r"""Benchmark for evaluating agents in environments."""
 
     def __init__(
-        self, config_paths: Optional[str] = None, eval_remote=False
+        self, config_paths: Optional[str] = None, eval_remote: bool = False
     ) -> None:
         r"""..
 
@@ -38,7 +40,7 @@ def __init__(
             self._env = Env(config=config_env)
 
     def remote_evaluate(
-        self, agent: Agent, num_episodes: Optional[int] = None
+        self, agent: "Agent", num_episodes: Optional[int] = None
     ):
         # The modules imported below are specific to habitat-challenge remote evaluation.
         # These modules are not part of the habitat-lab repository.
@@ -113,7 +115,9 @@ def remote_ep_over(stub):
 
         return avg_metrics
 
-    def local_evaluate(self, agent: Agent, num_episodes: Optional[int] = None):
+    def local_evaluate(
+        self, agent: "Agent", num_episodes: Optional[int] = None
+    ) -> Dict[str, float]:
         if num_episodes is None:
             num_episodes = len(self._env.episodes)
         else:
@@ -147,7 +151,7 @@ def local_evaluate(self, agent: Agent, num_episodes: Optional[int] = None):
         return avg_metrics
 
     def evaluate(
-        self, agent: Agent, num_episodes: Optional[int] = None
+        self, agent: "Agent", num_episodes: Optional[int] = None
     ) -> Dict[str, float]:
         r"""..
 

habitat/core/dataset.py

@@ -21,11 +21,14 @@
     Iterator,
     List,
     Optional,
+    Sequence,
     TypeVar,
+    Union,
 )
 
 import attr
 import numpy as np
+from numpy import ndarray
 
 from habitat.config import Config
 from habitat.core.utils import not_none_validator
@@ -58,7 +61,7 @@ class Episode:
     start_rotation: List[float] = attr.ib(
         default=None, validator=not_none_validator
     )
-    info: Optional[Dict[str, str]] = None
+    info: Optional[Dict[str, Any]] = None
     _shortest_path_cache: Any = attr.ib(init=False, default=None)
 
     def __getstate__(self):
@@ -101,8 +104,8 @@ def get_scenes_to_load(cls, config: Config) -> List[str]:
 
         :return: A list of scene names that would be loaded with the dataset
         """
-        assert cls.check_config_paths_exist(config)
-        dataset = cls(config)
+        assert cls.check_config_paths_exist(config)  # type: ignore[attr-defined]
+        dataset = cls(config)  # type: ignore[call-arg]
         return list(map(cls.scene_from_scene_path, dataset.scene_ids))
 
     @classmethod
@@ -279,7 +282,7 @@ def get_splits(
             self.num_episodes, num_episodes, replace=False
         )
         if collate_scene_ids:
-            scene_ids = {}
+            scene_ids: Dict[str, List[int]] = {}
             for rand_ind in rand_items:
                 scene = self.episodes[rand_ind].scene_id
                 if scene not in scene_ids:
@@ -334,7 +337,7 @@ class EpisodeIterator(Iterator):
 
     def __init__(
         self,
-        episodes: List[T],
+        episodes: Sequence[T],
         cycle: bool = True,
         shuffle: bool = False,
         group_by_scene: bool = True,
@@ -343,7 +346,7 @@ def __init__(
         num_episode_sample: int = -1,
         step_repetition_range: float = 0.2,
         seed: int = None,
-    ):
+    ) -> None:
         r"""..
 
         :param episodes: list of episodes.
@@ -375,7 +378,7 @@ def __init__(
                 episodes, num_episode_sample, replace=False
             )
 
-        self.episodes = episodes
+        self.episodes = list(episodes)
         self.cycle = cycle
         self.group_by_scene = group_by_scene
         self.shuffle = shuffle
@@ -391,17 +394,17 @@ def __init__(
 
         self._rep_count = -1  # 0 corresponds to first episode already returned
         self._step_count = 0
-        self._prev_scene_id = None
+        self._prev_scene_id: Optional[str] = None
 
         self._iterator = iter(self.episodes)
 
         self.step_repetition_range = step_repetition_range
         self._set_shuffle_intervals()
 
-    def __iter__(self):
+    def __iter__(self) -> "EpisodeIterator":
         return self
 
-    def __next__(self):
+    def __next__(self) -> Episode:
         r"""The main logic for handling how episodes will be iterated.
 
         :return: next episode.
@@ -459,7 +462,9 @@ def _shuffle(self) -> None:
 
         self._iterator = iter(episodes)
 
-    def _group_scenes(self, episodes):
+    def _group_scenes(
+        self, episodes: Union[Sequence[Episode], List[Episode], ndarray]
+    ) -> List[T]:
         r"""Internal method that groups episodes by scene
         Groups will be ordered by the order the first episode of a given
         scene is in the list of episodes
@@ -469,23 +474,23 @@ def _group_scenes(self, episodes):
         """
         assert self.group_by_scene
 
-        scene_sort_keys = {}
+        scene_sort_keys: Dict[str, int] = {}
         for e in episodes:
             if e.scene_id not in scene_sort_keys:
                 scene_sort_keys[e.scene_id] = len(scene_sort_keys)
 
-        return sorted(episodes, key=lambda e: scene_sort_keys[e.scene_id])
+        return sorted(episodes, key=lambda e: scene_sort_keys[e.scene_id])  # type: ignore[arg-type]
 
-    def step_taken(self):
+    def step_taken(self) -> None:
         self._step_count += 1
 
     @staticmethod
-    def _randomize_value(value, value_range):
+    def _randomize_value(value: int, value_range: float) -> int:
         return random.randint(
             int(value * (1 - value_range)), int(value * (1 + value_range))
         )
 
-    def _set_shuffle_intervals(self):
+    def _set_shuffle_intervals(self) -> None:
         if self.max_scene_repetition_episodes > 0:
             self._max_rep_episode = self.max_scene_repetition_episodes
         else:
@@ -498,7 +503,7 @@ def _set_shuffle_intervals(self):
         else:
             self._max_rep_step = None
 
-    def _forced_scene_switch_if(self):
+    def _forced_scene_switch_if(self) -> None:
         do_switch = False
         self._rep_count += 1
 

habitat/core/embodied_task.py

@@ -288,7 +288,7 @@ def reset(self, episode: Type[Episode]):
 
         return observations
 
-    def step(self, action: Union[int, Dict[str, Any]], episode: Type[Episode]):
+    def step(self, action: Dict[str, Any], episode: Type[Episode]):
         if "action_args" not in action or action["action_args"] is None:
             action["action_args"] = {}
         action_name = action["action"]

habitat/core/env.py

@@ -342,14 +342,14 @@ def habitat_env(self) -> Env:
     def episodes(self) -> List[Type[Episode]]:
         return self._env.episodes
 
-    @property
-    def current_episode(self) -> Type[Episode]:
-        return self._env.current_episode
-
     @episodes.setter
     def episodes(self, episodes: List[Type[Episode]]) -> None:
         self._env.episodes = episodes
 
+    @property
+    def current_episode(self) -> Type[Episode]:
+        return self._env.current_episode
+
     def reset(self) -> Observations:
         return self._env.reset()
 

habitat/core/logging.py

@@ -21,9 +21,9 @@ def __init__(
     ):
         super().__init__(name, level)
         if filename is not None:
-            handler = logging.FileHandler(filename, filemode)
+            handler = logging.FileHandler(filename, filemode)  # type:ignore
         else:
-            handler = logging.StreamHandler(stream)
+            handler = logging.StreamHandler(stream)  # type:ignore
         self._formatter = logging.Formatter(format, dateformat, style)
         handler.setFormatter(self._formatter)
         super().addHandler(handler)

habitat/core/registry.py

@@ -28,16 +28,29 @@
 """
 
 import collections
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Callable, DefaultDict, Optional, Type
 
 from habitat.core.utils import Singleton
 
+if TYPE_CHECKING:
+    from habitat.core.dataset import Dataset
+    from habitat.core.embodied_task import Measure
+    from habitat.core.simulator import ActionSpaceConfiguration, Sensor
+    from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim
+    from habitat.tasks.nav.nav import NavigationTask, SimulatorTaskAction
+
 
 class Registry(metaclass=Singleton):
-    mapping = collections.defaultdict(dict)
+    mapping: DefaultDict[str, Any] = collections.defaultdict(dict)
 
     @classmethod
-    def _register_impl(cls, _type, to_register, name, assert_type=None):
+    def _register_impl(
+        cls,
+        _type: str,
+        to_register: Optional[Any],
+        name: str,
+        assert_type: Optional[Type] = None,
+    ) -> Callable:
         def wrap(to_register):
             if assert_type is not None:
                 assert issubclass(
@@ -87,7 +100,7 @@ class MyTask(EmbodiedTask):
 
     @classmethod
     def register_simulator(
-        cls, to_register=None, *, name: Optional[str] = None
+        cls, to_register: None = None, *, name: Optional[str] = None
     ):
         r"""Register a simulator to registry with key :p:`name`
 
@@ -193,35 +206,37 @@ def register_action_space_configuration(
         )
 
     @classmethod
-    def _get_impl(cls, _type, name):
+    def _get_impl(cls, _type: str, name: str) -> Type:
         return cls.mapping[_type].get(name, None)
 
     @classmethod
-    def get_task(cls, name):
+    def get_task(cls, name: str) -> Type["NavigationTask"]:
         return cls._get_impl("task", name)
 
     @classmethod
-    def get_task_action(cls, name):
+    def get_task_action(cls, name: str) -> Type["SimulatorTaskAction"]:
         return cls._get_impl("task_action", name)
 
     @classmethod
-    def get_simulator(cls, name):
+    def get_simulator(cls, name: str) -> Type["HabitatSim"]:
         return cls._get_impl("sim", name)
 
     @classmethod
-    def get_sensor(cls, name):
+    def get_sensor(cls, name: str) -> Type["Sensor"]:
         return cls._get_impl("sensor", name)
 
     @classmethod
-    def get_measure(cls, name):
+    def get_measure(cls, name: str) -> Type["Measure"]:
         return cls._get_impl("measure", name)
 
     @classmethod
-    def get_dataset(cls, name):
+    def get_dataset(cls, name: str) -> Type["Dataset"]:
         return cls._get_impl("dataset", name)
 
     @classmethod
-    def get_action_space_configuration(cls, name):
+    def get_action_space_configuration(
+        cls, name: str
+    ) -> Type["ActionSpaceConfiguration"]:
         return cls._get_impl("action_space_config", name)
 
 

habitat/core/simulator.py

@@ -6,7 +6,7 @@
 
 from collections import OrderedDict
 from enum import Enum
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Union
 
 import attr
 from gym import Space
@@ -259,8 +259,8 @@ def reconfigure(self, config: Config) -> None:
 
     def geodesic_distance(
         self,
-        position_a: List[float],
-        position_b: Union[List[float], List[List[float]]],
+        position_a: Sequence[float],
+        position_b: Union[Sequence[float], Sequence[Sequence[float]]],
         episode: Optional[Episode] = None,
     ) -> float:
         r"""Calculates geodesic distance between two points.