fix imprecise resource type-hinting (#534)

ihmeuw · Nov 13, 2024 · 8491afb · 8491afb
1 parent 6c454ce
commit 8491afb
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 12 deletions.
diff --git a/src/vivarium/component.py b/src/vivarium/component.py
@@ -32,8 +32,7 @@
     from vivarium.framework.event import Event
     from vivarium.framework.lookup import LookupTable
     from vivarium.framework.population import PopulationView, SimulantData
-    from vivarium.framework.randomness import RandomnessStream
-    from vivarium.framework.values import Pipeline
+    from vivarium.framework.resource import Resource
     from vivarium.types import LookupTableData
 
 DEFAULT_EVENT_PRIORITY = 5
@@ -239,7 +238,7 @@ def columns_required(self) -> Optional[List[str]]:
     @property
     def initialization_requirements(
         self,
-    ) -> list[str | Pipeline | RandomnessStream]:
+    ) -> list[str | Resource]:
         """A list containing the columns, pipelines, and randomness streams
         required by this component's simulant initializer."""
         return []

diff --git a/src/vivarium/examples/disease_model/risk.py b/src/vivarium/examples/disease_model/risk.py
@@ -9,8 +9,7 @@
 
 if TYPE_CHECKING:
     from vivarium.framework.engine import Builder
-    from vivarium.framework.randomness import RandomnessStream
-    from vivarium.framework.values import Pipeline
+    from vivarium.framework.resource import Resource
 
 
 class Risk(Component):
@@ -33,7 +32,7 @@ def columns_created(self) -> List[str]:
         return [self.propensity_column]
 
     @property
-    def initialization_requirements(self) -> list[str | Pipeline | RandomnessStream]:
+    def initialization_requirements(self) -> list[str | Resource]:
         return [self.randomness]
 
     #####################

diff --git a/src/vivarium/framework/state_machine.py b/src/vivarium/framework/state_machine.py
@@ -22,8 +22,7 @@
     from vivarium.framework.engine import Builder
     from vivarium.framework.event import Event
     from vivarium.framework.population import PopulationView, SimulantData
-    from vivarium.framework.randomness import RandomnessStream
-    from vivarium.framework.values import Pipeline
+    from vivarium.framework.resource import Resource
     from vivarium.types import ClockTime, LookupTableData
 
 
@@ -495,7 +494,7 @@ def columns_created(self) -> List[str]:
     @property
     def initialization_requirements(
         self,
-    ) -> list[str | Pipeline | RandomnessStream]:
+    ) -> list[str | Resource]:
         return [self.randomness]
 
     #####################

diff --git a/tests/helpers.py b/tests/helpers.py
@@ -9,8 +9,7 @@
 from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
-from vivarium.framework.randomness import RandomnessStream
-from vivarium.framework.values import Pipeline
+from vivarium.framework.resource import Resource
 
 
 class MockComponentA(Observer):
@@ -247,7 +246,7 @@ def columns_created(self) -> List[str]:
         return ["test_column_4"]
 
     @property
-    def initialization_requirements(self) -> list[str | Pipeline | RandomnessStream]:
+    def initialization_requirements(self) -> list[str | Resource]:
         return ["test_column_2", self.pipeline, self.randomness]
 
     def setup(self, builder: Builder) -> None: