From d41ddcb3f3aaf560ee0550f52732b4120d725c22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Dallaire=20C=C3=B4t=C3=A9?=
 <110583667+0mdc@users.noreply.github.com>
Date: Thu, 6 Jun 2024 12:33:12 -0400
Subject: [PATCH] Add missing typing information to articulated agent manager.
 (#1977)

---
 .../rearrange/articulated_agent_manager.py    | 20 ++++++++++---------
 habitat-lab/habitat/tasks/rearrange/utils.py  |  2 +-
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py b/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py
index dbc5e3e95f..6fcee1d61f 100644
--- a/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py
+++ b/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py
@@ -3,7 +3,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional
 
 import magnum as mn
 import numpy as np
@@ -33,6 +33,8 @@
 if TYPE_CHECKING:
     from omegaconf import DictConfig
 
+    from habitat_sim.simulator import Simulator
+
 
 @dataclass
 class ArticulatedAgentData:
@@ -67,11 +69,11 @@ class ArticulatedAgentManager:
     Handles creating, updating and managing all agent instances.
     """
 
-    def __init__(self, cfg, sim):
+    def __init__(self, cfg: "DictConfig", sim: "Simulator"):
         self._sim = sim
-        self._all_agent_data = []
+        self._all_agent_data: List[ArticulatedAgentData] = []
         self._is_pb_installed = is_pb_installed()
-        self.agent_names = cfg.agents
+        self.agent_names: Dict[str, Any] = cfg.agents
 
         for agent_name in cfg.agents_order:
             agent_cfg = cfg.agents[agent_name]
@@ -80,10 +82,10 @@ def __init__(self, cfg, sim):
             agent = agent_cls(agent_cfg, sim)
             grasp_managers = []
             for grasp_manager_id in range(agent_cfg.grasp_managers):
-                graps_mgr = RearrangeGraspManager(
+                grasp_mgr = RearrangeGraspManager(
                     sim, cfg, agent, grasp_manager_id
                 )
-                grasp_managers.append(graps_mgr)
+                grasp_managers.append(grasp_mgr)
 
             if len(cfg.agents) > 1:
                 # Prefix sensors if there is more than 1 agent in the scene.
@@ -170,7 +172,7 @@ def post_obj_load_reconfigure(self):
             )
             agent_data.articulated_agent.reset()
 
-            # consume a fixed position from SIMUALTOR.agent_0 if configured
+            # Consume a fixed position from SIMULATOR.agent_0 if configured
             if agent_data.cfg.is_set_start_state:
                 agent_data.articulated_agent.base_pos = mn.Vector3(
                     agent_data.cfg.start_position
@@ -180,14 +182,14 @@ def post_obj_load_reconfigure(self):
                     mn.Vector3(agent_rot[:3]), agent_rot[3]
                 )
 
-    def __getitem__(self, key: int):
+    def __getitem__(self, key: int) -> ArticulatedAgentData:
         """
         Fetches the agent data at the agent index.
         """
 
         return self._all_agent_data[key]
 
-    def __len__(self):
+    def __len__(self) -> int:
         """
         The number of agents.
         """
diff --git a/habitat-lab/habitat/tasks/rearrange/utils.py b/habitat-lab/habitat/tasks/rearrange/utils.py
index bf4558fa63..9745456476 100644
--- a/habitat-lab/habitat/tasks/rearrange/utils.py
+++ b/habitat-lab/habitat/tasks/rearrange/utils.py
@@ -354,7 +354,7 @@ def batch_transform_point(
     p = None
 
 
-def is_pb_installed():
+def is_pb_installed() -> bool:
     return p is not None