fix mypy errors in framework/resource.py (#490)

CHANGELOG.rst

@@ -1,6 +1,7 @@
 **3.0.8 - 10/03/24**
 
-  = Fix mypy errors: vivarium/framework/artifact/hdf.py
+  - Fix mypy errors: vivarium/framework/resource.py
+  - Fix mypy errors: vivarium/framework/artifact/hdf.py
 
 **3.0.7 - 09/25/24**
 

pyproject.toml

@@ -61,7 +61,6 @@ exclude = [
     'src/vivarium/framework/population/population_view.py',
     'src/vivarium/framework/randomness/manager.py',
     'src/vivarium/framework/randomness/stream.py',
-    'src/vivarium/framework/resource.py',
     'src/vivarium/framework/results/context.py',
     'src/vivarium/framework/results/interface.py',
     'src/vivarium/framework/results/manager.py',

setup.py

@@ -44,6 +44,7 @@
         "dill",
         # Type stubs
         "pandas-stubs",
+        "networkx-stubs",
     ]
 
     setup_requires = ["setuptools_scm"]

src/vivarium/framework/resource.py

@@ -1,4 +1,3 @@
-# mypy: ignore-errors
 """
 ===================
 Resource Management
@@ -18,14 +17,19 @@
 
 """
 
-from types import MethodType
-from typing import Any, Callable, Iterable, List
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
 
 import networkx as nx
 
 from vivarium.exceptions import VivariumError
 from vivarium.manager import Interface, Manager
 
+if TYPE_CHECKING:
+    from vivarium.framework.engine import Builder
+
 
 class ResourceError(VivariumError):
     """Error raised when a dependency requirement is violated."""
@@ -59,9 +63,9 @@ class ResourceGroup:
     def __init__(
         self,
         resource_type: str,
-        resource_names: List[str],
-        producer: Callable,
-        dependencies: List[str],
+        resource_names: list[str],
+        producer: Any,
+        dependencies: list[str],
     ):
         self._resource_type = resource_type
         self._resource_names = resource_names
@@ -77,7 +81,7 @@ def type(self) -> str:
         return self._resource_type
 
     @property
-    def names(self) -> List[str]:
+    def names(self) -> list[str]:
         """The long names (including type) of all resources in this group."""
         return [f"{self._resource_type}.{name}" for name in self._resource_names]
 
@@ -87,11 +91,11 @@ def producer(self) -> Any:
         return self._producer
 
     @property
-    def dependencies(self) -> List[str]:
+    def dependencies(self) -> list[str]:
         """The long names (including type) of dependencies for this group."""
         return self._dependencies
 
-    def __iter__(self) -> Iterable[str]:
+    def __iter__(self) -> Iterator[str]:
         return iter(self.names)
 
     def __repr__(self) -> str:
@@ -106,20 +110,20 @@ def __str__(self) -> str:
 class ResourceManager(Manager):
     """Manages all the resources needed for population initialization."""
 
-    def __init__(self):
-        # This will be a dict with string keys representing the the resource
-        # and the resource group they belong to. This is a one to many mapping
-        # as some resource groups contain many resources.
-        self._resource_group_map = {}
-        # null producers are those that don't produce any resources externally
-        # but still consume other resources (i.e., have dependencies) - these
-        # are only pop initializers as of 9/26/2019. Tracker is here to assign
-        # them unique ids.
+    def __init__(self) -> None:
+        self._resource_group_map: dict[str, ResourceGroup] = {}
+        """Resource - resource group mapping. This will be a dict with string keys 
+        representing the the resource and the resource group they belong to. This 
+        is a one to many mapping as some resource groups contain many resources."""
         self._null_producer_count = 0
-        # Attribute used for lazy (but cached) graph initialization.
-        self._graph = None
-        # Attribute used for lazy (but cached) graph topo sort.
-        self._sorted_nodes = None
+        """Null producer counter. Null producers are those that don't produce any 
+        resources externally but still consume other resources (i.e., have 
+        dependencies) - these are only pop initializers as of 9/26/2019. Tracker 
+        is here to assign them unique ids."""
+        self._graph: nx.DiGraph | None = None
+        """Attribute used for lazy (but cached) graph initialization."""
+        self._sorted_nodes: list[ResourceGroup] | None = None
+        """Attribute used for lazy (but cached) graph topo sort."""
 
     @property
     def name(self) -> str:
@@ -134,7 +138,7 @@ def graph(self) -> nx.DiGraph:
         return self._graph
 
     @property
-    def sorted_nodes(self):
+    def sorted_nodes(self) -> list[ResourceGroup]:
         """Returns a topological sort of the resource graph.
 
         Notes
@@ -144,24 +148,25 @@ def sorted_nodes(self):
         """
         if self._sorted_nodes is None:
             try:
-                self._sorted_nodes = list(nx.algorithms.topological_sort(self.graph))
+                self._sorted_nodes = list(nx.algorithms.topological_sort(self.graph))  # type: ignore[func-returns-value]
             except nx.NetworkXUnfeasible:
                 raise ResourceError(
                     f"The resource pool contains at least one cycle: "
                     f"{nx.find_cycle(self.graph)}."
                 )
         return self._sorted_nodes
 
-    def setup(self, builder):
+    def setup(self, builder: Builder) -> None:
         self.logger = builder.logging.get_logger(self.name)
 
+    # TODO [MIC-5380]: Refactor add_resources for better type hinting
     def add_resources(
         self,
         resource_type: str,
-        resource_names: List[str],
+        resource_names: list[str],
         producer: Any,
-        dependencies: List[str],
-    ):
+        dependencies: list[str],
+    ) -> None:
         """Adds managed resources to the resource pool.
 
         Parameters
@@ -206,9 +211,9 @@ def add_resources(
     def _get_resource_group(
         self,
         resource_type: str,
-        resource_names: List[str],
-        producer: MethodType,
-        dependencies: List[str],
+        resource_names: list[str],
+        producer: Any,
+        dependencies: list[str],
     ) -> ResourceGroup:
         """Packages resource information into a resource group.
 
@@ -250,7 +255,7 @@ def _to_graph(self) -> nx.DiGraph:
                 if dependency not in self._resource_group_map:
                     # Warn here because this sometimes happens naturally
                     # if observer components are missing from a simulation.
-                    self.logger.warning(
+                    self.logger.warning(  # type: ignore[no-untyped-call]
                         f"Resource {dependency} is not provided by any component but is needed to "
                         f"compute {resource_group}."
                     )
@@ -260,7 +265,7 @@ def _to_graph(self) -> nx.DiGraph:
 
         return resource_graph
 
-    def __iter__(self) -> Iterable[MethodType]:
+    def __iter__(self) -> Iterator[Any]:
         """Returns a dependency-sorted iterable of population initializers.
 
         We exclude all non-initializer dependencies. They were necessary in
@@ -275,7 +280,7 @@ def __iter__(self) -> Iterable[MethodType]:
             ]
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         out = {}
         for resource_group in set(self._resource_group_map.values()):
             produced = ", ".join(resource_group)
@@ -307,9 +312,9 @@ def __init__(self, manager: ResourceManager):
     def add_resources(
         self,
         resource_type: str,
-        resource_names: List[str],
+        resource_names: list[str],
         producer: Any,
-        dependencies: List[str],
+        dependencies: list[str],
     ) -> None:
         """Adds managed resources to the resource pool.
 
@@ -335,7 +340,7 @@ def add_resources(
         """
         self._manager.add_resources(resource_type, resource_names, producer, dependencies)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         """Returns a dependency-sorted iterable of population initializers.
 
         We exclude all non-initializer dependencies. They were necessary in