Fix remaining types 2/2

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

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)
 

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]] = {}

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:
         ...
 

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",

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,

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."""
 

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/

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).

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