From 70df703510651c2814a1a8b9e91078102a824b3c Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 28 Oct 2021 04:50:26 +0200 Subject: [PATCH] Fix remaining types 2/2 --- openfisca_core/populations/population.py | 2 +- openfisca_core/simulations/simulation.py | 5 ++-- .../simulations/simulation_builder.py | 12 ++++++--- .../taxbenefitsystems/tax_benefit_system.py | 2 +- openfisca_core/typing/__init__.py | 3 +++ openfisca_core/typing/_protocols.py | 26 +++++++++++++++++-- openfisca_core/typing/_schemas.py | 11 ++++++++ openfisca_core/typing/_types.py | 7 +++-- openfisca_tasks/lint.mk | 4 +-- setup.cfg | 1 + 10 files changed, 59 insertions(+), 14 deletions(-) diff --git a/openfisca_core/populations/population.py b/openfisca_core/populations/population.py index 0caec1be14..fa13ef3070 100644 --- a/openfisca_core/populations/population.py +++ b/openfisca_core/populations/population.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional from openfisca_core.typing import ArrayLike import traceback diff --git a/openfisca_core/simulations/simulation.py b/openfisca_core/simulations/simulation.py index 24604ab505..c737260839 100644 --- a/openfisca_core/simulations/simulation.py +++ b/openfisca_core/simulations/simulation.py @@ -18,7 +18,6 @@ from openfisca_core.errors import CycleError, SpiralError from openfisca_core.indexed_enums import Enum, EnumArray from openfisca_core.periods import Period -from openfisca_core.populations import Population from openfisca_core.tracers import FullTracer, SimpleTracer, TracingParameterNodeAtInstant from openfisca_core.warnings import TempfileWarning @@ -99,7 +98,7 @@ def data_storage_dir(self): def calculate( self, variable_name: str, - period: Period, + period: Optional[Any], ) -> ArrayType[Any]: """Calculate ``variable_name`` for ``period``.""" @@ -439,7 +438,7 @@ def get_variable_population(self, variable_name): def get_population( self, plural: Optional[str] = None, - ) -> Optional[Population]: + ) -> Optional[PopulationProtocol]: return next((population for population in self.populations.values() if population.entity.plural == plural), None) diff --git a/openfisca_core/simulations/simulation_builder.py b/openfisca_core/simulations/simulation_builder.py index e6b7f51413..552fa2b278 100644 --- a/openfisca_core/simulations/simulation_builder.py +++ b/openfisca_core/simulations/simulation_builder.py @@ -1,8 +1,12 @@ from __future__ import annotations import typing -from typing import Any, Mapping, Optional -from openfisca_core.typing import ArrayType, TaxBenefitSystemProtocol +from typing import Any, Mapping, Optional, Sequence +from openfisca_core.typing import ( + ArrayType, + AxisSchema, + TaxBenefitSystemProtocol, + ) import copy import dpath @@ -19,6 +23,8 @@ class SimulationBuilder: + default_period: Optional[str] + def __init__(self) -> None: self.default_period = None # Simulation period used for variables when no period is defined self.persons_plural = None # Plural name for person entity in current tax and benefits system @@ -37,7 +43,7 @@ def __init__(self) -> None: self.variable_entities: typing.Dict[Variable.name, Entity] = {} - self.axes = [[]] + self.axes: Sequence[Sequence[AxisSchema]] = [[]] self.axes_entity_counts: typing.Dict[Entity.plural, int] = {} self.axes_entity_ids: typing.Dict[Entity.plural, typing.List[int]] = {} self.axes_memberships: typing.Dict[Entity.plural, typing.List[int]] = {} diff --git a/openfisca_core/taxbenefitsystems/tax_benefit_system.py b/openfisca_core/taxbenefitsystems/tax_benefit_system.py index 3c4b90c673..4d8d0bd555 100644 --- a/openfisca_core/taxbenefitsystems/tax_benefit_system.py +++ b/openfisca_core/taxbenefitsystems/tax_benefit_system.py @@ -293,7 +293,7 @@ def apply_reform(self, reform_path): def get_variable( self, variable_name: str, - check_existence: Literal[True], + check_existence: Literal[True] = ..., ) -> Variable: ... diff --git a/openfisca_core/typing/__init__.py b/openfisca_core/typing/__init__.py index 5f9a28d607..648a96ffc9 100644 --- a/openfisca_core/typing/__init__.py +++ b/openfisca_core/typing/__init__.py @@ -12,6 +12,7 @@ * :class:`.PeriodProtocol` * :class:`.PopulationProtocol` * :class:`.TaxBenefitSystemProtocol` + * :class:`.AxisSchema` * :class:`.FrameSchema` * :class:`.OptionsSchema` * :class:`.TestSchema` @@ -60,6 +61,7 @@ ) from ._schemas import ( # noqa: F401 + AxisSchema, FrameSchema, OptionsSchema, TestSchema, @@ -73,6 +75,7 @@ "PeriodProtocol", "PopulationProtocol", "TaxBenefitSystemProtocol", + "AxisSchema", "FrameSchema", "OptionsSchema", "TestSchema", diff --git a/openfisca_core/typing/_protocols.py b/openfisca_core/typing/_protocols.py index fd21f94d38..a5a2c0405d 100644 --- a/openfisca_core/typing/_protocols.py +++ b/openfisca_core/typing/_protocols.py @@ -1,7 +1,9 @@ +# pylint: disable=missing-function-docstring + from __future__ import annotations -from typing import Any, Mapping, Optional, Sequence, Set -from typing_extensions import Protocol +from typing import Any, Mapping, Optional, Sequence, Set, overload +from typing_extensions import Literal, Protocol from ._types import ArrayType import abc @@ -117,6 +119,10 @@ class PopulationProtocol(Protocol): members_position: ArrayType[int] members_role: ArrayType[RoleProtocol] + @abc.abstractmethod + def get_index(self, id: str) -> int: + ... + class RoleProtocol(Protocol): """Duck-type for roles. @@ -165,6 +171,22 @@ def entities_plural(self) -> Set[str]: def get_package_metadata(self) -> Mapping[str, str]: ... + @overload + def get_variable( + self, + variable_name: str, + check_existence: Literal[True] = ..., + ) -> VariableProtocol: + ... + + @overload + def get_variable( + self, + variable_name: str, + check_existence: bool = ..., + ) -> Optional[VariableProtocol]: + ... + @abc.abstractmethod def get_variable( self, diff --git a/openfisca_core/typing/_schemas.py b/openfisca_core/typing/_schemas.py index cf786cf54c..b9d7a086cf 100644 --- a/openfisca_core/typing/_schemas.py +++ b/openfisca_core/typing/_schemas.py @@ -6,6 +6,17 @@ from ._protocols import PeriodProtocol +class AxisSchema(TypedDict): + """Data-schema of axes.""" + + count: int + index: int + max: float + min: float + name: str + period: str + + class FrameSchema(TypedDict): """Data-schema of tracer stack frames.""" diff --git a/openfisca_core/typing/_types.py b/openfisca_core/typing/_types.py index 5cfef639c5..8441e2dd9a 100644 --- a/openfisca_core/typing/_types.py +++ b/openfisca_core/typing/_types.py @@ -34,10 +34,13 @@ Todo: * Refactor once numpy version >= 1.21 is used. -.. versionadded:: 35.5.0 +.. versionchanged:: 35.8.0 + Moved to :mod:`.openfisca_core.typing` .. versionchanged:: 35.6.0 - Moved to :mod:`.types` + Moved to ``openfisca_core.types`` + +.. versionadded:: 35.5.0 .. _mypy: https://mypy.readthedocs.io/en/stable/ diff --git a/openfisca_tasks/lint.mk b/openfisca_tasks/lint.mk index 75182cc53a..117da9e6c3 100644 --- a/openfisca_tasks/lint.mk +++ b/openfisca_tasks/lint.mk @@ -17,7 +17,7 @@ check-style: $(shell git ls-files "*.py") ## Run linters to check for syntax and style errors in the doc. lint-doc: \ lint-doc-commons \ - lint-doc-types \ + lint-doc-typing \ ; ## Run linters to check for syntax and style errors in the doc. @@ -43,7 +43,7 @@ check-types: lint-typing-strict: \ lint-typing-strict-commons \ lint-typing-strict-tools \ - lint-typing-strict-types \ + lint-typing-strict-typing \ ; ## Run static type checkers for type errors (strict). diff --git a/setup.cfg b/setup.cfg index e69600ed39..ab06fe8f12 100644 --- a/setup.cfg +++ b/setup.cfg @@ -15,6 +15,7 @@ hang-closing = true ignore = E128,E251,F403,F405,E501,RST301,W503,W504 in-place = true include-in-doctest = openfisca_core/commons openfisca_core/typing +per-file-ignores = openfisca_core/typing/_protocols.py:D102 rst-directives = attribute, deprecated, seealso, versionadded, versionchanged rst-roles = any, attr, class, data, exc, func, meth, obj strictness = short