typing: some small typing improvements to API/context modules

Author: Fizzadar

pyinfra/api/inventory.py

@@ -36,6 +36,7 @@ class Inventory:
     """
 
     state: "State"
+    groups: dict[str, list[Host]]
 
     @staticmethod
     def empty():
@@ -181,7 +182,7 @@ def len_activated_hosts(self) -> int:
         """
         return len(self.state.activated_hosts)
 
-    def get_host(self, name: str, default=NoHostError):
+    def get_host(self, name: str, default=NoHostError) -> Host:
         """
         Get a single host by name.
         """
@@ -192,9 +193,10 @@ def get_host(self, name: str, default=NoHostError):
         if default is NoHostError:
             raise NoHostError("No such host: {0}".format(name))
 
+        # TODO: remove default here?
         return default
 
-    def get_group(self, name: str, default=NoGroupError):
+    def get_group(self, name: str, default=NoGroupError) -> list[Host]:
         """
         Get a list of hosts belonging to a group.
         """
@@ -205,6 +207,7 @@ def get_group(self, name: str, default=NoGroupError):
         if default is NoGroupError:
             raise NoGroupError("No such group: {0}".format(name))
 
+        # TODO: remove default here?
         return default
 
     def get_data(self):

pyinfra/context.py

@@ -59,16 +59,22 @@ def __setattr__(self, key, value):
         if key in ("_container", "_base_cls"):
             return super().__setattr__(key, value)
 
-        if self._get_module() is None:
+        mod = self._get_module()
+        if mod is None:
             raise TypeError("Cannot assign to context base module")
-
-        return setattr(self._get_module(), key, value)
+        return setattr(mod, key, value)
 
     def __iter__(self):
-        return iter(self._get_module())
+        mod = self._get_module()
+        if mod is None:
+            raise ValueError("Context not set")
+        return iter(mod)
 
     def __len__(self):
-        return len(self._get_module())
+        mod = self._get_module()
+        if mod is None:
+            raise ValueError("Context not set")
+        return len(mod)
 
     @override
     def __eq__(self, other):
@@ -105,6 +111,9 @@ def isset(self):
     @contextmanager
     def use(self, module):
         old_module = self.get()
+        if old_module is module:
+            yield  # if we're double-setting, nothing to do
+            return
         self.set(module)
         yield
         self.set(old_module)