From 1c299cbb79c7a4622a46cf101f8ff819c99c4c9d Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Tue, 17 Sep 2024 00:56:15 +0200 Subject: [PATCH] test(simulations): improve simulation builder doc --- .../simulations/_build_from_variables.py | 6 +- openfisca_core/simulations/_guards.py | 495 ++++++++++++++---- openfisca_core/simulations/simulation.py | 19 +- openfisca_core/simulations/types.py | 90 +++- openfisca_tasks/lint.mk | 7 +- setup.cfg | 2 +- setup.py | 1 + 7 files changed, 481 insertions(+), 139 deletions(-) diff --git a/openfisca_core/simulations/_build_from_variables.py b/openfisca_core/simulations/_build_from_variables.py index 152e73846..292a921d8 100644 --- a/openfisca_core/simulations/_build_from_variables.py +++ b/openfisca_core/simulations/_build_from_variables.py @@ -8,7 +8,7 @@ from openfisca_core import errors from ._build_default_simulation import _BuildDefaultSimulation -from ._guards import is_variable_dated +from ._guards import is_a_dated_value, is_a_pure_value from .simulation import Simulation from .types import Populations, TaxBenefitSystem, Variables @@ -142,7 +142,7 @@ def add_dated_values(self) -> Self: """ for variable, value in self.variables.items(): - if is_variable_dated(dated_variable := value): + if is_a_dated_value(dated_variable := value): for period, dated_value in dated_variable.items(): self.simulation.set_input(variable, period, dated_value) @@ -200,7 +200,7 @@ def add_undated_values(self) -> Self: """ for variable, value in self.variables.items(): - if not is_variable_dated(undated_value := value): + if is_a_pure_value(undated_value := value): if (period := self.default_period) is None: message = ( "Can't deal with type: expected object. Input " diff --git a/openfisca_core/simulations/_guards.py b/openfisca_core/simulations/_guards.py index 79e0c8584..404d026d2 100644 --- a/openfisca_core/simulations/_guards.py +++ b/openfisca_core/simulations/_guards.py @@ -1,82 +1,427 @@ -"""Type guards to help type narrowing simulation parameters.""" +"""Type guards to help type narrowing simulation parameters. + +Every calculation in a simulation requires an entity, a variable, a period, and +a value. However, the way users can specify these elements can vary. This +module provides type guards to help narrow down the type of simulation +parameters, to help both readability and maintainability. + +For example, the following is a perfectly valid, albeit complex, way to specify +a simulation's parameters:: + + .. code-block:: python + + params = { + "axes": [ + [ + { + "count": 2, + "max": 3000, + "min": 0, + "name": + "rent", + "period": "2018-11" + } + ] + ], + "households": { + "housea": { + "parents": ["Alicia", "Javier"] + }, + "houseb": { + "parents": ["Tom"] + }, + }, + "persons": { + "Alicia": { + "salary": { + "2018-11": 0 + } + }, + "Javier": {}, + "Tom": {} + }, + } + +""" from __future__ import annotations -from typing import Iterable +from collections.abc import Iterable from typing_extensions import TypeGuard +import pydantic + +from openfisca_core import periods + from .types import ( Axes, - DatedVariable, + Axis, + DatedValue, FullySpecifiedEntities, ImplicitGroupEntities, Params, - UndatedVariable, + ParamsWithAxes, + PureValue, + Roles, Variables, ) +#: Pydantic type adapter to extract information from axes. +adapter = pydantic.TypeAdapter(Axis) -def are_entities_fully_specified( +#: Field schema for axes. +axis_schema = adapter.core_schema + +#: Required fields. +axis_required = [ + key for key, value in axis_schema["fields"].items() if value["required"] +] + + +def is_a_pure_value( + value: object, +) -> TypeGuard[PureValue]: + """Check if an input value is undated. + + The most atomic elements of a simulation are pure values. They can be + either scalars or vectors. For example:: + + .. code-block:: python + + 1.5 + True + [1000, 2000] + + Args: + value(object): A value. + + Returns: + bool: True if the value is undated. + + Examples: + >>> value = 2000 + >>> is_a_pure_value(value) + True + + >>> value = [2000, 3000] + >>> is_a_pure_value(value) + True + + >>> value = {"2000": 2000} + >>> is_a_pure_value(value) + False + + >>> value = {"2018-W01": [2000, 3000]} + >>> is_a_pure_value(value) + False + + >>> value = {"123": 123} + >>> is_a_pure_value(value) + False + + """ + + return not isinstance(value, dict) + + +def is_a_dated_value( + value: object, +) -> TypeGuard[DatedValue]: + """Check if an input value is dated. + + Pure values are associated with the simulation's period behind the scenes. + However, some calculations require different values for variables for + different periods. In such a case, users can specify dated values:: + + .. code-block:: python + + {"2018-01": 2000} + {"2018-W01": [2000, 3000]} + {"2018-W01-1": 2000, "2018-W01-2": [3000, 4000]} + + Args: + value(object): A value. + + Returns: + bool: True if the value is dated. + + Examples: + >>> value = 2000 + >>> is_a_dated_value(value) + False + + >>> value = [2000, 3000] + >>> is_a_dated_value(value) + False + + >>> value = {"2000": 2000} + >>> is_a_dated_value(value) + True + + >>> value = {"2018-W01": [2000, 3000]} + >>> is_a_dated_value(value) + True + + >>> value = {"123": 123} + >>> is_a_dated_value(value) + False + + """ + + if not isinstance(value, dict): + return False + + try: + return all(periods.period(key) for key in value.keys()) + + except ValueError: + return False + + +def are_variables( + value: object, +) -> TypeGuard[Variables]: + """Check if an input value is a map of variables. + + In a simulation, every value has to be associated with a variable. As with + values, variables cannot be inferred from the context. Users have to + explicitly specify them. For example:: + + .. code-block:: python + + {"salary": 2000} + {"taxes": {"2018-W01-1": [123, 234]}} + {"taxes": {"2018-W01-1": [123, 234]}, "salary": 123} + + Args: + value(object): A value. + + Returns: + bool: True if the value is a map of variables. + + Examples: + >>> value = 2000 + >>> are_variables(value) + False + + >>> value = [2000, 3000] + >>> are_variables(value) + False + + >>> value = {"2000": 2000} + >>> are_variables(value) + False + + >>> value = {"2018-W01": [2000, 3000]} + >>> are_variables(value) + False + + >>> value = {"salary": 123} + >>> are_variables(value) + True + + >>> value = {"taxes": {"2018-W01-1": [123, 234]}} + >>> are_variables(value) + True + + >>> value = {"taxes": {"2018-W01-1": [123, 234]}, "salary": 123} + >>> are_variables(value) + True + + """ + + if is_a_pure_value(value): + return False + + if is_a_dated_value(value): + return False + + return True + + +def are_roles( + value: object, +) -> TypeGuard[Roles]: + """Check if an input value is a map of roles. + + In a simulation, there are cases where we need to calculate things for + group entities, for example, households. In such cases, some calculations + require that we specify certain roles. For example:: + + .. code-block:: python + + {"principal": "Alicia"} + {"parents": ["Alicia", "Javier"]} + + Args: + value(object): A value. + + Returns: + bool: True if the value is a map of roles. + + Examples: + >>> value = "parent" + >>> are_roles(value) + False + + >>> value = ["dad", "son"] + >>> are_roles(value) + False + + >>> value = {"2018-W01": [2000, 3000]} + >>> are_roles(value) + False + + >>> value = {"salary": 123} + >>> are_roles(value) + False + + >>> value = {"principal": "Alicia"} + >>> are_roles(value) + True + + >>> value = {"kids": ["Alicia", "Javier"]} + >>> are_roles(value) + True + + >>> value = {"principal": "Alicia", "kids": ["Tom"]} + >>> are_roles(value) + True + + """ + + if not isinstance(value, dict): + return False + + for role_key, role_id in value.items(): + if not isinstance(role_key, str): + return False + + if not isinstance(role_id, (Iterable, str)): + return False + + if isinstance(role_id, Iterable): + for role in role_id: + if not isinstance(role, str): + return False + + return True + + +def are_axes(value: object) -> TypeGuard[Axes]: + """Check if the given params are axes. + + Axis expansion is a feature that allows users to parametrise some + dimensions in order to create and to evaluate a range of values for others. + + Args: + value(object): Simulation parameters. + + Returns: + bool: True if the params are axes. + + Examples: + >>> value = { + ... "persons": {"Javier": { "salary": { "2018-11": 2000}}}, + ... "households": {"household": {"parents": ["Javier"]}}, + ... "axes": [[{"count": 1, "max": 1, "min": 1, "name": "household"}]] + ... } + >>> are_axes(value) + False + + >>> value = { + ... "axes": [[{"count": 1, "max": 1, "min": 1, "name": "household"}]] + ... } + >>> are_axes(value) + False + + >>> value = [[{"a": 1, "b": 1, "c": 1}]] + >>> are_axes(value) + False + + >>> value = [[{"count": 1, "max": 1, "min": 1, "name": "household"}]] + >>> are_axes(value) + True + + """ + + if not isinstance(value, (list, tuple)): + return False + + (inner,) = value + + if not isinstance(inner, (list, tuple)): + return False + + return all(key in axis_required for key in inner[0].keys()) + + +def are_entities_specified( params: Params, items: Iterable[str] -) -> TypeGuard[FullySpecifiedEntities]: - """Check if the params contain fully specified entities. +) -> TypeGuard[Variables]: + """Check if the params contains entities at all. Args: params(Params): Simulation parameters. - items(Iterable[str]): List of entities in plural form. + items(Iterable[str]): List of variables. Returns: - bool: True if the params contain fully specified entities. + bool: True if the params does not contain variables at the root level. Examples: - >>> entities = {"persons", "households"} + >>> variables = {"salary"} >>> params = { - ... "axes": [ - ... [{"count": 2, "max": 3000, "min": 0, "name": "rent", "period": "2018-11"}] - ... ], - ... "households": { - ... "housea": {"parents": ["Alicia", "Javier"]}, - ... "houseb": {"parents": ["Tom"]}, - ... }, - ... "persons": {"Alicia": {"salary": {"2018-11": 0}}, "Javier": {}, "Tom": {}}, + ... "persons": {"Javier": { "salary": { "2018-11": 2000}}}, + ... "households": {"household": {"parents": ["Javier"]}}, + ... "axes": [[{"count": 1, "max": 1, "min": 1, "name": "household"}]] ... } - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) True >>> params = { ... "persons": {"Javier": {"salary": {"2018-11": [2000, 3000]}}} ... } - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) True >>> params = { ... "persons": {"Javier": {"salary": {"2018-11": 2000}}} ... } - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) True >>> params = {"household": {"parents": ["Javier"]}} - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) + True + + >>> params = {"salary": {"2016-10": [12000, 13000]}} + + >>> are_entities_specified(params, variables) False >>> params = {"salary": {"2016-10": 12000}} - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) + False + + >>> params = {"salary": [12000, 13000]} + + >>> are_entities_specified(params, variables) False >>> params = {"salary": 12000} - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) False >>> params = {} - >>> are_entities_fully_specified(params, entities) + >>> are_entities_specified(params, variables) False """ @@ -84,7 +429,7 @@ def are_entities_fully_specified( if not params: return False - return all(key in items for key in params.keys() if key != "axes") + return not any(key in items for key in params.keys()) def are_entities_short_form( @@ -165,72 +510,67 @@ def are_entities_short_form( return not not set(params).intersection(items) -def are_entities_specified( +def are_entities_fully_specified( params: Params, items: Iterable[str] -) -> TypeGuard[Variables]: - """Check if the params contains entities at all. +) -> TypeGuard[FullySpecifiedEntities]: + """Check if the params contain fully specified entities. Args: params(Params): Simulation parameters. - items(Iterable[str]): List of variables. + items(Iterable[str]): List of entities in plural form. Returns: - bool: True if the params does not contain variables at the root level. + bool: True if the params contain fully specified entities. Examples: - >>> variables = {"salary"} + >>> entities = {"persons", "households"} >>> params = { - ... "persons": {"Javier": { "salary": { "2018-11": 2000}}}, - ... "households": {"household": {"parents": ["Javier"]}}, - ... "axes": [[{"count": 1, "max": 1, "min": 1, "name": "household"}]] + ... "axes": [ + ... [{"count": 2, "max": 3000, "min": 0, "name": "rent", "period": "2018-11"}] + ... ], + ... "households": { + ... "housea": {"parents": ["Alicia", "Javier"]}, + ... "houseb": {"parents": ["Tom"]}, + ... }, + ... "persons": {"Alicia": {"salary": {"2018-11": 0}}, "Javier": {}, "Tom": {}}, ... } - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) True >>> params = { ... "persons": {"Javier": {"salary": {"2018-11": [2000, 3000]}}} ... } - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) True >>> params = { ... "persons": {"Javier": {"salary": {"2018-11": 2000}}} ... } - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) True >>> params = {"household": {"parents": ["Javier"]}} - >>> are_entities_specified(params, variables) - True - - >>> params = {"salary": {"2016-10": [12000, 13000]}} - - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) False >>> params = {"salary": {"2016-10": 12000}} - >>> are_entities_specified(params, variables) - False - - >>> params = {"salary": [12000, 13000]} - - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) False >>> params = {"salary": 12000} - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) False >>> params = {} - >>> are_entities_specified(params, variables) + >>> are_entities_fully_specified(params, entities) False """ @@ -238,67 +578,36 @@ def are_entities_specified( if not params: return False - return not any(key in items for key in params.keys()) + return all(key in items for key in params.keys() if key != "axes") -def has_axes(params: Params) -> TypeGuard[Axes]: +def has_axes(value: object) -> TypeGuard[ParamsWithAxes]: """Check if the params contains axes. Args: - params(Params): Simulation parameters. + value(object): Simulation parameters. Returns: bool: True if the params contain axes. Examples: - >>> params = { + >>> value = { ... "persons": {"Javier": { "salary": { "2018-11": 2000}}}, ... "households": {"household": {"parents": ["Javier"]}}, ... "axes": [[{"count": 1, "max": 1, "min": 1, "name": "household"}]] ... } - - >>> has_axes(params) + >>> has_axes(value) True - >>> params = { + >>> value = { ... "persons": {"Javier": {"salary": {"2018-11": [2000, 3000]}}} ... } - - >>> has_axes(params) + >>> has_axes(value) False """ - return params.get("axes", None) is not None - - -def is_variable_dated( - variable: DatedVariable | UndatedVariable, -) -> TypeGuard[DatedVariable]: - """Check if the variable is dated. - - Args: - variable(DatedVariable | UndatedVariable): A variable. - - Returns: - bool: True if the variable is dated. - - Examples: - >>> variable = {"2018-11": [2000, 3000]} - - >>> is_variable_dated(variable) - True - - >>> variable = {"2018-11": 2000} - - >>> is_variable_dated(variable) - True - - >>> variable = 2000 - - >>> is_variable_dated(variable) - False - - """ + if not isinstance(value, dict): + return False - return isinstance(variable, dict) + return value.get("axes", None) is not None diff --git a/openfisca_core/simulations/simulation.py b/openfisca_core/simulations/simulation.py index 136b960fd..7cd2fe4c6 100644 --- a/openfisca_core/simulations/simulation.py +++ b/openfisca_core/simulations/simulation.py @@ -17,6 +17,7 @@ Populations, TaxBenefitSystem, Variable, + VariableName, ) @@ -100,7 +101,7 @@ def data_storage_dir(self): # ----- Calculation methods ----- # - def calculate(self, variable_name: str, period): + def calculate(self, variable_name: VariableName, period): """Calculate ``variable_name`` for ``period``.""" if period is not None and not isinstance(period, periods.Period): @@ -117,7 +118,7 @@ def calculate(self, variable_name: str, period): self.tracer.record_calculation_end() self.purge_cache_of_invalid_values() - def _calculate(self, variable_name: str, period: periods.Period): + def _calculate(self, variable_name: VariableName, period: periods.Period): """ Calculate the variable ``variable_name`` for the period ``period``, using the variable formula if it exists. @@ -169,7 +170,7 @@ def purge_cache_of_invalid_values(self): holder.delete_arrays(_period) self.invalidated_caches = set() - def calculate_add(self, variable_name: str, period): + def calculate_add(self, variable_name: VariableName, period): variable: Optional[Variable] variable = self.tax_benefit_system.get_variable( @@ -207,7 +208,7 @@ def calculate_add(self, variable_name: str, period): for sub_period in period.get_subperiods(variable.definition_period) ) - def calculate_divide(self, variable_name: str, period): + def calculate_divide(self, variable_name: VariableName, period): variable: Optional[Variable] variable = self.tax_benefit_system.get_variable( @@ -284,7 +285,7 @@ def calculate_divide(self, variable_name: str, period): return self.calculate(variable_name, calculation_period) / denominator - def calculate_output(self, variable_name: str, period): + def calculate_output(self, variable_name: VariableName, period): """ Calculate the value of a variable using the ``calculate_output`` attribute of the variable. """ @@ -434,7 +435,7 @@ def invalidate_spiral_variables(self, variable: str): # ----- Methods to access stored values ----- # - def get_array(self, variable_name: str, period): + def get_array(self, variable_name: VariableName, period): """ Return the value of ``variable_name`` for ``period``, if this value is alreay in the cache (if it has been set as an input or previously calculated). @@ -444,7 +445,7 @@ def get_array(self, variable_name: str, period): period = periods.period(period) return self.get_holder(variable_name).get_array(period) - def get_holder(self, variable_name: str): + def get_holder(self, variable_name: VariableName): """Get the holder associated with the variable.""" return self.get_variable_population(variable_name).get_holder(variable_name) @@ -507,7 +508,7 @@ def get_known_periods(self, variable): """ return self.get_holder(variable).get_known_periods() - def set_input(self, variable_name: str, period, value): + def set_input(self, variable_name: VariableName, period, value): """ Set a variable's value for a given period @@ -538,7 +539,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) -> GroupPopulation: + def get_variable_population(self, variable_name: VariableName) -> GroupPopulation: variable: Optional[Variable] variable = self.tax_benefit_system.get_variable( diff --git a/openfisca_core/simulations/types.py b/openfisca_core/simulations/types.py index cce6d12b8..68dd7dc96 100644 --- a/openfisca_core/simulations/types.py +++ b/openfisca_core/simulations/types.py @@ -3,8 +3,8 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Sequence -from typing import NewType, Protocol, TypeVar, TypedDict, Union -from typing_extensions import NotRequired, Required, TypeAlias +from typing import Literal, NewType, Protocol, TypeVar, Union +from typing_extensions import NotRequired, TypeAlias, TypedDict import datetime @@ -24,9 +24,29 @@ V = TypeVar("V", covariant=True) # New types. -PeriodStr = NewType("PeriodStr", str) + +#: Literally "axes". +AxesKey = Literal["axes"] + +#: For example "Juan". +EntityId = NewType("EntityId", int) + +#: For example "person". EntityKey = NewType("EntityKey", str) + +#: For example "persons". EntityPlural = NewType("EntityPlural", str) + +#: For example "2023-12". +PeriodStr = NewType("PeriodStr", str) + +#: For example "principal". +RoleKey = NewType("RoleKey", str) + +#: For example "parents". +RolePlural = NewType("RolePlural", str) + +#: For example "salary". VariableName = NewType("VariableName", str) # Type aliases. @@ -37,6 +57,9 @@ #: Type Alias for a numpy Array. Array: TypeAlias = t.Array +#: Type alias for a role identifier. +RoleId: TypeAlias = EntityId + # Entities @@ -88,6 +111,10 @@ def set_input( # Periods +class Instant(t.Instant, Protocol): + ... + + class Period(t.Period, Protocol): ... @@ -119,35 +146,38 @@ def nb_persons(self, __role: Role[G] | None = ...) -> int: #: Dictionary with axes parameters per variable. InputBuffer: TypeAlias = dict[VariableName, dict[PeriodStr, Array]] -#: Dictionary with entity/population key/pais. +#: Dictionary with entity/population key/pairs. Populations: TypeAlias = dict[EntityKey, GroupPopulation] #: Dictionary with single entity count per group entity. EntityCounts: TypeAlias = dict[EntityPlural, int] #: Dictionary with a list of single entities per group entity. -EntityIds: TypeAlias = dict[EntityPlural, Iterable[int]] +EntityIds: TypeAlias = dict[EntityPlural, Iterable[EntityId]] #: Dictionary with a list of members per group entity. Memberships: TypeAlias = dict[EntityPlural, Iterable[int]] #: Dictionary with a list of roles per group entity. -EntityRoles: TypeAlias = dict[EntityPlural, Iterable[int]] +EntityRoles: TypeAlias = dict[EntityPlural, Iterable[RoleKey]] #: Dictionary with a map between variables and entities. VariableEntity: TypeAlias = dict[VariableName, CoreEntity] -#: Type alias for a simulation dictionary defining the roles. -Roles: TypeAlias = dict[str, Union[str, Iterable[str]]] - -#: Type alias for a simulation dictionary with undated variables. -UndatedVariable: TypeAlias = dict[str, object] +#: Type alias for a simulation dictionary with undated variable values. +PureValue: TypeAlias = Union[object, Sequence[object]] -#: Type alias for a simulation dictionary with dated variables. -DatedVariable: TypeAlias = dict[str, UndatedVariable] +#: Type alias for a simulation dictionary with dated variable values. +DatedValue: TypeAlias = dict[PeriodStr, PureValue] #: Type alias for a simulation dictionary with abbreviated entities. -Variables: TypeAlias = dict[str, Union[UndatedVariable, DatedVariable]] +Variables: TypeAlias = dict[VariableName, Union[PureValue, DatedValue]] + +#: Type alias for a simulation dictionary defining the roles. +Roles: TypeAlias = Union[dict[RoleKey, RoleId], dict[RolePlural, Iterable[RoleId]]] + +#: Type alias for a simulation dictionary with axes parameters. +Axes: TypeAlias = Iterable[Iterable["Axis"]] #: Type alias for a simulation with fully specified single entities. SingleEntities: TypeAlias = dict[str, dict[str, Variables]] @@ -161,28 +191,25 @@ def nb_persons(self, __role: Role[G] | None = ...) -> int: #: Type alias for a simulation dictionary with fully specified entities. FullySpecifiedEntities: TypeAlias = Union[SingleEntities, GroupEntities] -#: Type alias for a simulation dictionary with axes parameters. -Axes: TypeAlias = dict[str, Iterable[Iterable["Axis"]]] - #: Type alias for a simulation dictionary without axes parameters. ParamsWithoutAxes: TypeAlias = Union[ Variables, ImplicitGroupEntities, FullySpecifiedEntities ] #: Type alias for a simulation dictionary with axes parameters. -ParamsWithAxes: TypeAlias = Union[Axes, ParamsWithoutAxes] +ParamsWithAxes: TypeAlias = Union[dict[AxesKey, Axes], ParamsWithoutAxes] #: Type alias for a simulation dictionary with all the possible scenarios. Params: TypeAlias = ParamsWithAxes -class Axis(TypedDict, total=False): - count: Required[int] +class Axis(TypedDict): + count: int + max: float + min: float index: NotRequired[int] - max: Required[float] - min: Required[float] - name: Required[str] - period: NotRequired[str | int] + name: EntityKey + period: NotRequired[Union[str, int]] class Simulation(t.Simulation, Protocol): @@ -205,15 +232,15 @@ def person_entity(self, person_entity: SingleEntity) -> None: def variables(self) -> dict[str, V]: ... - def entities_by_singular(self) -> dict[str, CoreEntity]: + def entities_by_singular(self) -> dict[EntityKey, CoreEntity]: ... - def entities_plural(self) -> Iterable[str]: + def entities_plural(self) -> Iterable[EntityPlural]: ... def get_variable( self, - __variable_name: str, + __variable_name: VariableName, check_existence: bool = ..., ) -> Variable[T] | None: ... @@ -235,3 +262,12 @@ class Variable(t.Variable, Protocol[T]): def default_array(self, __array_size: int) -> t.Array[T]: ... + + def get_formula( + self, __period: Instant | Period | PeriodStr | Int + ) -> Formula | None: + ... + + +class Formula(t.Formula, Protocol): + ... diff --git a/openfisca_tasks/lint.mk b/openfisca_tasks/lint.mk index c0b026502..b58d61699 100644 --- a/openfisca_tasks/lint.mk +++ b/openfisca_tasks/lint.mk @@ -43,12 +43,7 @@ check-types: openfisca_core/commons \ openfisca_core/entities \ openfisca_core/projectors \ - openfisca_core/simulations/_build_default_simulation.py \ - openfisca_core/simulations/_build_from_variables.py \ - openfisca_core/simulations/_guards.py \ - openfisca_core/simulations/helpers.py \ - openfisca_core/simulations/simulation.py \ - openfisca_core/simulations/types.py \ + openfisca_core/simulations \ openfisca_core/types.py @$(call print_pass,$@:) diff --git a/setup.cfg b/setup.cfg index cc850c06a..5c47d1cd8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -75,7 +75,7 @@ ignore_missing_imports = true implicit_reexport = false install_types = true non_interactive = true -plugins = numpy.typing.mypy_plugin +plugins = numpy.typing.mypy_plugin, pydantic.mypy pretty = true python_version = 3.9 strict = false diff --git a/setup.py b/setup.py index 92971f40b..46a5e0194 100644 --- a/setup.py +++ b/setup.py @@ -33,6 +33,7 @@ "numpy >=1.24.2, <1.25", "pendulum >=2.1.2, <3.0.0", "psutil >=5.9.4, <6.0", + "pydantic >=2.9.1, <3.0", "pytest >=7.2.2, <8.0", "sortedcontainers >=2.4.0, <3.0", "typing_extensions >=4.5.0, <5.0",