diff --git a/openfisca_core/projectors/__init__.py b/openfisca_core/projectors/__init__.py index 28776e3cf..958251082 100644 --- a/openfisca_core/projectors/__init__.py +++ b/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", ] diff --git a/openfisca_core/projectors/helpers.py b/openfisca_core/projectors/helpers.py index b3b7e6f2d..9c666eb6b 100644 --- a/openfisca_core/projectors/helpers.py +++ b/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) diff --git a/openfisca_core/projectors/types.py b/openfisca_core/projectors/types.py new file mode 100644 index 000000000..a090c4813 --- /dev/null +++ b/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]: + ... diff --git a/openfisca_core/projectors/typing.py b/openfisca_core/projectors/typing.py deleted file mode 100644 index 186f90e30..000000000 --- a/openfisca_core/projectors/typing.py +++ /dev/null @@ -1,32 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Protocol - -from openfisca_core.types import GroupEntity, SingleEntity - - -class Population(Protocol): - @property - def entity(self) -> SingleEntity: - ... - - @property - def simulation(self) -> Simulation: - ... - - -class GroupPopulation(Protocol): - @property - def entity(self) -> GroupEntity: - ... - - @property - def simulation(self) -> Simulation: - ... - - -class Simulation(Protocol): - @property - def populations(self) -> Mapping[str, Population | GroupPopulation]: - ... diff --git a/openfisca_core/simulations/simulation.py b/openfisca_core/simulations/simulation.py index 93becda96..c4525525d 100644 --- a/openfisca_core/simulations/simulation.py +++ b/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 diff --git a/openfisca_core/types.py b/openfisca_core/types.py index b34a55543..ba3a2f2c0 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -113,7 +113,11 @@ def unit(self) -> Any: # Populations -class Population(Protocol): +class CorePopulation(Protocol): + ... + + +class SinglePopulation(CorePopulation, Protocol): entity: Any @abc.abstractmethod @@ -121,6 +125,10 @@ def get_holder(self, variable_name: Any) -> Any: ... +class GroupPopulation(CorePopulation, Protocol): + ... + + # Simulations @@ -168,7 +176,7 @@ class Formula(Protocol): @abc.abstractmethod def __call__( self, - population: Population, + population: GroupPopulation, instant: Instant, params: Params, ) -> Array[Any]: diff --git a/openfisca_tasks/lint.mk b/openfisca_tasks/lint.mk index 445abba10..2ba76a07b 100644 --- a/openfisca_tasks/lint.mk +++ b/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,$@:)