fix(types): projectors

openfisca_core/entities/types.py

@@ -1,34 +1,30 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Protocol, TypedDict
+from typing import Protocol
+from typing_extensions import TypedDict
 
-from openfisca_core import types
+from openfisca_core import types as t
 
 # Entities
 
 
-class CoreEntity(types.CoreEntity, Protocol):
+class CoreEntity(t.CoreEntity, Protocol):
     ...
 
 
-class SingleEntity(types.SingleEntity, Protocol):
+class SingleEntity(t.SingleEntity, Protocol):
     key: str
     plural: str | None
 
 
-class GroupEntity(types.GroupEntity, Protocol):
+class GroupEntity(t.GroupEntity, Protocol):
     ...
 
 
-class Role(types.Role, Protocol):
-    max: int | None
+class Role(t.Role, Protocol):
     subroles: Iterable[Role] | None
 
-    @property
-    def key(self) -> str:
-        ...
-
 
 class RoleParams(TypedDict, total=False):
     key: str

openfisca_core/projectors/__init__.py

@@ -21,7 +21,7 @@
 #
 # See: https://www.python.org/dev/peps/pep-0008/#imports
 
-from . import typing
+from . import types
 from .entity_to_person_projector import EntityToPersonProjector
 from .first_person_to_entity_projector import FirstPersonToEntityProjector
 from .helpers import get_projector_from_shortcut, projectable
@@ -35,5 +35,5 @@
     "projectable",
     "Projector",
     "UniqueRoleToEntityProjector",
-    "typing",
+    "types",
 ]

openfisca_core/projectors/helpers.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
-
-from openfisca_core.types import GroupEntity, Role, SingleEntity
+from collections.abc import Iterable, Mapping
 
 from openfisca_core import entities, projectors
 
-from .typing import GroupPopulation, Population
+from .types import GroupEntity, GroupPopulation, Role, SingleEntity, SinglePopulation
 
 
 def projectable(function):
@@ -19,7 +17,7 @@ def projectable(function):
 
 
 def get_projector_from_shortcut(
-    population: Population | GroupPopulation,
+    population: SinglePopulation | GroupPopulation,
     shortcut: str,
     parent: projectors.Projector | None = None,
 ) -> projectors.Projector | None:
@@ -46,7 +44,7 @@ def get_projector_from_shortcut(
     of a specific Simulation and TaxBenefitSystem.
 
     Args:
-        population (Population | GroupPopulation): Where to project from.
+        population (SinglePopulation | GroupPopulation): Where to project from.
         shortcut (str): Where to project to.
         parent: ???
 
@@ -114,7 +112,7 @@ def get_projector_from_shortcut(
 
     if isinstance(entity, entities.Entity):
         populations: Mapping[
-            str, Population | GroupPopulation
+            str, SinglePopulation | GroupPopulation
         ] = population.simulation.populations
 
         if shortcut not in populations.keys():
@@ -126,7 +124,8 @@ def get_projector_from_shortcut(
         return projectors.FirstPersonToEntityProjector(population, parent)
 
     if isinstance(entity, entities.GroupEntity):
-        role: Role | None = entities.find_role(entity.roles, shortcut, total=1)
+        roles: Iterable[Role] = entity.roles
+        role: Role | None = entities.find_role(roles, shortcut, total=1)
 
         if role is not None:
             return projectors.UniqueRoleToEntityProjector(population, role, parent)

openfisca_core/projectors/types.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Protocol
+
+from openfisca_core import types as t
+
+# Entities
+
+
+class SingleEntity(t.SingleEntity, Protocol):
+    ...
+
+
+class GroupEntity(t.GroupEntity, Protocol):
+    ...
+
+
+class Role(t.Role, Protocol):
+    ...
+
+
+# Populations
+
+
+class SinglePopulation(t.SinglePopulation, Protocol):
+    @property
+    def entity(self) -> t.SingleEntity:
+        ...
+
+    @property
+    def simulation(self) -> Simulation:
+        ...
+
+
+class GroupPopulation(t.GroupPopulation, Protocol):
+    @property
+    def entity(self) -> t.GroupEntity:
+        ...
+
+    @property
+    def simulation(self) -> Simulation:
+        ...
+
+
+# Simulations
+
+
+class Simulation(t.Simulation, Protocol):
+    @property
+    def populations(self) -> Mapping[str, SinglePopulation | GroupPopulation]:
+        ...

openfisca_core/simulations/simulation.py

@@ -2,7 +2,7 @@
 
 from typing import Dict, Mapping, NamedTuple, Optional, Set
 
-from openfisca_core.types import Population, TaxBenefitSystem, Variable
+from openfisca_core.types import SinglePopulation, TaxBenefitSystem, Variable
 
 import tempfile
 import warnings
@@ -19,13 +19,13 @@ class Simulation:
     """
 
     tax_benefit_system: TaxBenefitSystem
-    populations: Dict[str, Population]
+    populations: Dict[str, SinglePopulation]
     invalidated_caches: Set[Cache]
 
     def __init__(
         self,
         tax_benefit_system: TaxBenefitSystem,
-        populations: Mapping[str, Population],
+        populations: Mapping[str, SinglePopulation],
     ):
         """
         This constructor is reserved for internal use; see :any:`SimulationBuilder`,
@@ -531,7 +531,7 @@ def set_input(self, variable_name: str, period, value):
             return
         self.get_holder(variable_name).set_input(period, value)
 
-    def get_variable_population(self, variable_name: str) -> Population:
+    def get_variable_population(self, variable_name: str) -> SinglePopulation:
         variable: Optional[Variable]
 
         variable = self.tax_benefit_system.get_variable(
@@ -543,7 +543,9 @@ def get_variable_population(self, variable_name: str) -> Population:
 
         return self.populations[variable.entity.key]
 
-    def get_population(self, plural: Optional[str] = None) -> Optional[Population]:
+    def get_population(
+        self, plural: Optional[str] = None
+    ) -> Optional[SinglePopulation]:
         return next(
             (
                 population
@@ -556,7 +558,7 @@ def get_population(self, plural: Optional[str] = None) -> Optional[Population]:
     def get_entity(
         self,
         plural: Optional[str] = None,
-    ) -> Optional[Population]:
+    ) -> Optional[SinglePopulation]:
         population = self.get_population(plural)
         return population and population.entity
 

openfisca_core/types.py

@@ -61,8 +61,13 @@ class GroupEntity(CoreEntity, Protocol):
 
 class Role(Protocol):
     entity: Any
+    max: int | None
     subroles: Any
 
+    @property
+    def key(self) -> str:
+        ...
+
 
 # Holders
 
@@ -108,14 +113,22 @@ def unit(self) -> Any:
 # Populations
 
 
-class Population(Protocol):
+class CorePopulation(Protocol):
+    ...
+
+
+class SinglePopulation(CorePopulation, Protocol):
     entity: Any
 
     @abc.abstractmethod
     def get_holder(self, variable_name: Any) -> Any:
         ...
 
 
+class GroupPopulation(CorePopulation, Protocol):
+    ...
+
+
 # Simulations
 
 
@@ -163,7 +176,7 @@ class Formula(Protocol):
     @abc.abstractmethod
     def __call__(
         self,
-        population: Population,
+        population: GroupPopulation,
         instant: Instant,
         params: Params,
     ) -> Array[Any]:

openfisca_tasks/lint.mk

@@ -42,6 +42,7 @@ check-types:
 	@mypy \
 		openfisca_core/commons \
 		openfisca_core/entities \
+		openfisca_core/projectors \
 		openfisca_core/types.py
 	@$(call print_pass,$@:)