From e24d2076d6cac0bf4ffa655a850c31bfa79907a4 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Tue, 17 Sep 2024 13:37:31 +0200 Subject: [PATCH 01/20] chore(deps): update pendulum --- openfisca_core/periods/period_.py | 8 ++++---- openfisca_tasks/lint.mk | 1 + setup.py | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/openfisca_core/periods/period_.py b/openfisca_core/periods/period_.py index 11a7b671b..6dfb25873 100644 --- a/openfisca_core/periods/period_.py +++ b/openfisca_core/periods/period_.py @@ -353,14 +353,14 @@ def size_in_weeks(self): if self.unit == DateUnit.YEAR: start = self.start.date cease = start.add(years=self.size) - delta = pendulum.period(start, cease) - return delta.as_interval().weeks + delta = pendulum.interval(start, cease) + return delta.in_weeks() if self.unit == DateUnit.MONTH: start = self.start.date cease = start.add(months=self.size) - delta = pendulum.period(start, cease) - return delta.as_interval().weeks + delta = pendulum.interval(start, cease) + return delta.in_weeks() if self.unit == DateUnit.WEEK: return self.size diff --git a/openfisca_tasks/lint.mk b/openfisca_tasks/lint.mk index 445abba10..7f90aa5d5 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/periods \ openfisca_core/types.py @$(call print_pass,$@:) diff --git a/setup.py b/setup.py index cca107bee..6707ce9c6 100644 --- a/setup.py +++ b/setup.py @@ -31,7 +31,7 @@ "dpath >=2.1.4, <3.0", "numexpr >=2.8.4, <3.0", "numpy >=1.24.2, <1.25", - "pendulum >=2.1.2, <3.0.0", + "pendulum >=3.0.0, <4.0.0", "psutil >=5.9.4, <6.0", "pytest >=7.2.2, <8.0", "sortedcontainers >=2.4.0, <3.0", From 711b42bf4407d2a0675711ec1a6a68c43d30e2a4 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 00:28:10 +0200 Subject: [PATCH 02/20] refactor(variables): remove rubn-time type check --- openfisca_core/periods/__init__.py | 49 +++++++++++++++++++++------- openfisca_core/periods/config.py | 9 ----- openfisca_core/periods/types.py | 0 openfisca_core/variables/variable.py | 9 +++-- 4 files changed, 42 insertions(+), 25 deletions(-) create mode 100644 openfisca_core/periods/types.py diff --git a/openfisca_core/periods/__init__.py b/openfisca_core/periods/__init__.py index 4669c7ff4..ca23dbf76 100644 --- a/openfisca_core/periods/__init__.py +++ b/openfisca_core/periods/__init__.py @@ -21,20 +21,14 @@ # # See: https://www.python.org/dev/peps/pep-0008/#imports -from .config import ( # noqa: F401 - DAY, - ETERNITY, +from .config import ( INSTANT_PATTERN, - MONTH, - WEEK, - WEEKDAY, - YEAR, date_by_instant_cache, str_by_instant_cache, year_or_month_or_day_re, ) -from .date_unit import DateUnit # noqa: F401 -from .helpers import ( # noqa: F401 +from .date_unit import DateUnit +from .helpers import ( instant, instant_date, key_period_size, @@ -42,5 +36,38 @@ unit_weight, unit_weights, ) -from .instant_ import Instant # noqa: F401 -from .period_ import Period # noqa: F401 +from .instant_ import Instant +from .period_ import Period + +WEEKDAY = DateUnit.WEEKDAY +WEEK = DateUnit.WEEK +DAY = DateUnit.DAY +MONTH = DateUnit.MONTH +YEAR = DateUnit.YEAR +ETERNITY = DateUnit.ETERNITY +ISOFORMAT = DateUnit.isoformat +ISOCALENDAR = DateUnit.isocalendar + +__all__ = [ + "INSTANT_PATTERN", + "date_by_instant_cache", + "str_by_instant_cache", + "year_or_month_or_day_re", + "DateUnit", + "instant", + "instant_date", + "key_period_size", + "period", + "unit_weight", + "unit_weights", + "Instant", + "Period", + "WEEKDAY", + "WEEK", + "DAY", + "MONTH", + "YEAR", + "ETERNITY", + "ISOFORMAT", + "ISOCALENDAR", +] diff --git a/openfisca_core/periods/config.py b/openfisca_core/periods/config.py index 17807160e..afbfbd9d0 100644 --- a/openfisca_core/periods/config.py +++ b/openfisca_core/periods/config.py @@ -1,14 +1,5 @@ import re -from .date_unit import DateUnit - -WEEKDAY = DateUnit.WEEKDAY -WEEK = DateUnit.WEEK -DAY = DateUnit.DAY -MONTH = DateUnit.MONTH -YEAR = DateUnit.YEAR -ETERNITY = DateUnit.ETERNITY - # Matches "2015", "2015-01", "2015-01-01" # Does not match "2015-13", "2015-12-32" INSTANT_PATTERN = re.compile( diff --git a/openfisca_core/periods/types.py b/openfisca_core/periods/types.py new file mode 100644 index 000000000..e69de29bb diff --git a/openfisca_core/variables/variable.py b/openfisca_core/variables/variable.py index e70c0d05d..1b073a1bd 100644 --- a/openfisca_core/variables/variable.py +++ b/openfisca_core/variables/variable.py @@ -2,8 +2,6 @@ from typing import Optional, Union -from openfisca_core.types import Formula, Instant - import datetime import re import textwrap @@ -12,6 +10,7 @@ import sortedcontainers from openfisca_core import periods, tools +from openfisca_core import types as t from openfisca_core.entities import Entity, GroupEntity from openfisca_core.indexed_enums import Enum, EnumArray from openfisca_core.periods import DateUnit, Period @@ -374,8 +373,8 @@ def get_introspection_data(cls): def get_formula( self, - period: Union[Instant, Period, str, int] = None, - ) -> Optional[Formula]: + period: Union[t.Instant, t.Period, str, int] = None, + ) -> Optional[t.Formula]: """Returns the formula to compute the variable at the given period. If no period is given and the variable has several formulas, the method @@ -389,7 +388,7 @@ def get_formula( """ - instant: Optional[Instant] + instant: Optional[t.Instant] if not self.formulas: return None From b2fd70b840b8c37e9879e7cb662588d3c30ca397 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 00:28:39 +0200 Subject: [PATCH 03/20] refactor(populations): remove run-time type check --- openfisca_core/populations/population.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfisca_core/populations/population.py b/openfisca_core/populations/population.py index e3ef6b209..4cdd7abcf 100644 --- a/openfisca_core/populations/population.py +++ b/openfisca_core/populations/population.py @@ -85,7 +85,7 @@ def check_period_validity( variable_name: str, period: Optional[Union[int, str, Period]], ) -> None: - if isinstance(period, (int, str, Period)): + if isinstance(period, (int, str, periods.Period)): return None stack = traceback.extract_stack() From d0b1b53bcd1e621aab205b681db8414f246ec93f Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 00:30:21 +0200 Subject: [PATCH 04/20] refactor(types): do not use abstract classes --- openfisca_core/types.py | 50 ++++++++++++++++++----------------------- 1 file changed, 22 insertions(+), 28 deletions(-) diff --git a/openfisca_core/types.py b/openfisca_core/types.py index b34a55543..fe9ea8dee 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -1,13 +1,10 @@ from __future__ import annotations -import typing_extensions -from collections.abc import Sequence +from collections.abc import Iterable, Sequence, Sized from numpy.typing import NDArray -from typing import Any, TypeVar +from typing import Any, TypeVar, Union from typing_extensions import Protocol, TypeAlias -import abc - import numpy N = TypeVar("N", bound=numpy.generic, covariant=True) @@ -26,6 +23,10 @@ #: Type variable representing a value. A = TypeVar("A", covariant=True) +#: Generic type vars. +T_cov = TypeVar("T_cov", covariant=True) +T_con = TypeVar("T_con", contravariant=True) + # Entities @@ -34,15 +35,12 @@ class CoreEntity(Protocol): key: Any plural: Any - @abc.abstractmethod def check_role_validity(self, role: Any) -> None: ... - @abc.abstractmethod def check_variable_defined_for_entity(self, variable_name: Any) -> None: ... - @abc.abstractmethod def get_variable( self, variable_name: Any, @@ -73,11 +71,9 @@ def key(self) -> str: class Holder(Protocol): - @abc.abstractmethod def clone(self, population: Any) -> Holder: ... - @abc.abstractmethod def get_memory_usage(self) -> Any: ... @@ -85,7 +81,6 @@ def get_memory_usage(self) -> Any: # Parameters -@typing_extensions.runtime_checkable class ParameterNodeAtInstant(Protocol): ... @@ -93,20 +88,27 @@ class ParameterNodeAtInstant(Protocol): # Periods -class Instant(Protocol): - ... +class Container(Protocol[T_con]): + def __contains__(self, item: T_con, /) -> bool: + ... -@typing_extensions.runtime_checkable -class Period(Protocol): - @property - @abc.abstractmethod - def start(self) -> Any: +class Indexable(Protocol[T_cov]): + def __getitem__(self, index: int, /) -> T_cov: ... + +class DateUnit(Container[str], Protocol): + ... + + +class Instant(Indexable[int], Iterable[int], Sized, Protocol): + ... + + +class Period(Indexable[Union[DateUnit, Instant, int]], Protocol): @property - @abc.abstractmethod - def unit(self) -> Any: + def unit(self) -> DateUnit: ... @@ -116,7 +118,6 @@ def unit(self) -> Any: class Population(Protocol): entity: Any - @abc.abstractmethod def get_holder(self, variable_name: Any) -> Any: ... @@ -125,19 +126,15 @@ def get_holder(self, variable_name: Any) -> Any: class Simulation(Protocol): - @abc.abstractmethod def calculate(self, variable_name: Any, period: Any) -> Any: ... - @abc.abstractmethod def calculate_add(self, variable_name: Any, period: Any) -> Any: ... - @abc.abstractmethod def calculate_divide(self, variable_name: Any, period: Any) -> Any: ... - @abc.abstractmethod def get_population(self, plural: Any | None) -> Any: ... @@ -148,7 +145,6 @@ def get_population(self, plural: Any | None) -> Any: class TaxBenefitSystem(Protocol): person_entity: Any - @abc.abstractmethod def get_variable( self, variable_name: Any, @@ -165,7 +161,6 @@ class Variable(Protocol): class Formula(Protocol): - @abc.abstractmethod def __call__( self, population: Population, @@ -176,6 +171,5 @@ def __call__( class Params(Protocol): - @abc.abstractmethod def __call__(self, instant: Instant) -> ParameterNodeAtInstant: ... From 787f583f18c8bf1e1ef3912c0a7410650cc1f57e Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 00:33:04 +0200 Subject: [PATCH 05/20] doc(periods): add types --- openfisca_core/periods/__init__.py | 2 + openfisca_core/periods/_parsers.py | 5 +- openfisca_core/periods/config.py | 8 +- openfisca_core/periods/date_unit.py | 6 +- openfisca_core/periods/helpers.py | 68 +++++----- openfisca_core/periods/instant_.py | 41 +++--- openfisca_core/periods/period_.py | 124 +++++++++--------- .../periods/tests/helpers/test_helpers.py | 14 +- .../periods/tests/helpers/test_period.py | 6 +- openfisca_core/periods/types.py | 60 +++++++++ 10 files changed, 217 insertions(+), 117 deletions(-) diff --git a/openfisca_core/periods/__init__.py b/openfisca_core/periods/__init__.py index ca23dbf76..1f2e68aca 100644 --- a/openfisca_core/periods/__init__.py +++ b/openfisca_core/periods/__init__.py @@ -21,6 +21,7 @@ # # See: https://www.python.org/dev/peps/pep-0008/#imports +from . import types from .config import ( INSTANT_PATTERN, date_by_instant_cache, @@ -70,4 +71,5 @@ "ETERNITY", "ISOFORMAT", "ISOCALENDAR", + "types", ] diff --git a/openfisca_core/periods/_parsers.py b/openfisca_core/periods/_parsers.py index 64b207783..0342cc71e 100644 --- a/openfisca_core/periods/_parsers.py +++ b/openfisca_core/periods/_parsers.py @@ -6,6 +6,7 @@ from pendulum.datetime import Date from pendulum.parsing import ParserError +from . import types as t from .date_unit import DateUnit from .instant_ import Instant from .period_ import Period @@ -13,7 +14,7 @@ invalid_week = re.compile(r".*(W[1-9]|W[1-9]-[0-9]|W[0-5][0-9]-0)$") -def _parse_period(value: str) -> Optional[Period]: +def _parse_period(value: str) -> Optional[t.Period]: """Parses ISO format/calendar periods. Such as "2012" or "2015-03". @@ -57,7 +58,7 @@ def _parse_period(value: str) -> Optional[Period]: return Period((unit, instant, 1)) -def _parse_unit(value: str) -> DateUnit: +def _parse_unit(value: str) -> t.DateUnit: """Determine the date unit of a date string. Args: diff --git a/openfisca_core/periods/config.py b/openfisca_core/periods/config.py index afbfbd9d0..3b4c0a7eb 100644 --- a/openfisca_core/periods/config.py +++ b/openfisca_core/periods/config.py @@ -1,13 +1,17 @@ import re +from pendulum.datetime import Date + +from . import types as t + # Matches "2015", "2015-01", "2015-01-01" # Does not match "2015-13", "2015-12-32" INSTANT_PATTERN = re.compile( r"^\d{4}(-(0[1-9]|1[012]))?(-(0[1-9]|1[012])-(0[1-9]|[12][0-9]|3[01]))?$" ) -date_by_instant_cache: dict = {} -str_by_instant_cache: dict = {} +date_by_instant_cache: dict[t.Instant, Date] = {} +str_by_instant_cache: dict[t.Instant, t.InstantStr] = {} year_or_month_or_day_re = re.compile( r"(18|19|20)\d{2}(-(0?[1-9]|1[0-2])(-([0-2]?\d|3[0-1]))?)?$" ) diff --git a/openfisca_core/periods/date_unit.py b/openfisca_core/periods/date_unit.py index a81321149..a808b1941 100644 --- a/openfisca_core/periods/date_unit.py +++ b/openfisca_core/periods/date_unit.py @@ -4,10 +4,12 @@ from strenum import StrEnum +from . import types as t + class DateUnitMeta(EnumMeta): @property - def isoformat(self) -> tuple[DateUnit, ...]: + def isoformat(self) -> tuple[t.DateUnit, ...]: """Creates a :obj:`tuple` of ``key`` with isoformat items. Returns: @@ -28,7 +30,7 @@ def isoformat(self) -> tuple[DateUnit, ...]: return DateUnit.DAY, DateUnit.MONTH, DateUnit.YEAR @property - def isocalendar(self) -> tuple[DateUnit, ...]: + def isocalendar(self) -> tuple[t.DateUnit, ...]: """Creates a :obj:`tuple` of ``key`` with isocalendar items. Returns: diff --git a/openfisca_core/periods/helpers.py b/openfisca_core/periods/helpers.py index 2ce4e0cd3..ce27c54d0 100644 --- a/openfisca_core/periods/helpers.py +++ b/openfisca_core/periods/helpers.py @@ -1,4 +1,6 @@ -from typing import NoReturn, Optional +from __future__ import annotations + +from typing import NoReturn, Optional, overload import datetime import os @@ -7,12 +9,23 @@ from pendulum.parsing import ParserError from . import _parsers, config +from . import types as t from .date_unit import DateUnit from .instant_ import Instant from .period_ import Period -def instant(instant) -> Optional[Instant]: +@overload +def instant(instant: None) -> None: + ... + + +@overload +def instant(instant: object) -> Instant: + ... + + +def instant(instant: object | None) -> t.Instant | None: """Build a new instant, aka a triple of integers (year, month, day). Args: @@ -49,6 +62,8 @@ def instant(instant) -> Optional[Instant]: """ + result: t.Instant | tuple[int, ...] + if instant is None: return None if isinstance(instant, Instant): @@ -58,27 +73,28 @@ def instant(instant) -> Optional[Instant]: raise ValueError( f"'{instant}' is not a valid instant. Instants are described using the 'YYYY-MM-DD' format, for instance '2015-06-15'." ) - instant = Instant(int(fragment) for fragment in instant.split("-", 2)[:3]) + result = Instant(int(fragment) for fragment in instant.split("-", 2)[:3]) elif isinstance(instant, datetime.date): - instant = Instant((instant.year, instant.month, instant.day)) + result = Instant((instant.year, instant.month, instant.day)) elif isinstance(instant, int): - instant = (instant,) + result = (instant,) elif isinstance(instant, list): assert 1 <= len(instant) <= 3 - instant = tuple(instant) + result = tuple(instant) elif isinstance(instant, Period): - instant = instant.start + result = instant.start else: assert isinstance(instant, tuple), instant assert 1 <= len(instant) <= 3 - if len(instant) == 1: - return Instant((instant[0], 1, 1)) - if len(instant) == 2: - return Instant((instant[0], instant[1], 1)) - return Instant(instant) + result = instant + if len(result) == 1: + return Instant((result[0], 1, 1)) + if len(result) == 2: + return Instant((result[0], result[1], 1)) + return Instant(result) -def instant_date(instant: Optional[Instant]) -> Optional[datetime.date]: +def instant_date(instant: Optional[t.Instant]) -> Optional[datetime.date]: """Returns the date representation of an :class:`.Instant`. Args: @@ -105,7 +121,7 @@ def instant_date(instant: Optional[Instant]) -> Optional[datetime.date]: return instant_date -def period(value) -> Period: +def period(value: object) -> t.Period: """Build a new period, aka a triple (unit, start_instant, size). Args: @@ -125,7 +141,7 @@ def period(value) -> Period: Period((, Instant((2021, 1, 1)), 1)) >>> period(DateUnit.ETERNITY) - Period((, Instant((1, 1, 1)), inf)) + Period((, Instant((1, 1, 1)), 0)) >>> period(2021) Period((, Instant((2021, 1, 1)), 1)) @@ -164,15 +180,9 @@ def period(value) -> Period: return Period((DateUnit.DAY, instant(value), 1)) # We return an "eternity-period", for example - # ``, inf))>``. + # ``, 0))>``. if str(value).lower() == DateUnit.ETERNITY: - return Period( - ( - DateUnit.ETERNITY, - instant(datetime.date.min), - float("inf"), - ) - ) + return Period.eternity() # For example ``2021`` gives # ``, 1))>``. @@ -181,7 +191,7 @@ def period(value) -> Period: # Up to this point, if ``value`` is not a :obj:`str`, we desist. if not isinstance(value, str): - _raise_error(value) + _raise_error(str(value)) # There can't be empty strings. if not value: @@ -267,7 +277,7 @@ def _raise_error(value: str) -> NoReturn: raise ValueError(message) -def key_period_size(period: Period) -> str: +def key_period_size(period: t.Period) -> str: """Define a key in order to sort periods by length. It uses two aspects: first, ``unit``, then, ``size``. @@ -291,12 +301,10 @@ def key_period_size(period: Period) -> str: """ - unit, start, size = period - - return f"{unit_weight(unit)}_{size}" + return f"{unit_weight(period.unit)}_{period.size}" -def unit_weights() -> dict[str, int]: +def unit_weights() -> dict[t.DateUnit, int]: """Assign weights to date units. Examples: @@ -315,7 +323,7 @@ def unit_weights() -> dict[str, int]: } -def unit_weight(unit: str) -> int: +def unit_weight(unit: t.DateUnit) -> int: """Retrieves a specific date unit weight. Examples: diff --git a/openfisca_core/periods/instant_.py b/openfisca_core/periods/instant_.py index 9d0893ba4..e68151c36 100644 --- a/openfisca_core/periods/instant_.py +++ b/openfisca_core/periods/instant_.py @@ -1,10 +1,14 @@ +from __future__ import annotations + import pendulum +from pendulum.datetime import Date from . import config +from . import types as t from .date_unit import DateUnit -class Instant(tuple): +class Instant(tuple[int, int, int]): """An instant in time (year, month, day). An :class:`.Instant` represents the most atomic and indivisible @@ -13,10 +17,6 @@ class Instant(tuple): Current implementation considers this unit to be a day, so :obj:`instants <.Instant>` can be thought of as "day dates". - Args: - (tuple(tuple(int, int, int))): - The ``year``, ``month``, and ``day``, accordingly. - Examples: >>> instant = Instant((2021, 9, 13)) @@ -78,35 +78,41 @@ class Instant(tuple): """ - def __repr__(self): + def __repr__(self) -> str: return f"{self.__class__.__name__}({super().__repr__()})" - def __str__(self): + def __str__(self) -> t.InstantStr: instant_str = config.str_by_instant_cache.get(self) if instant_str is None: - config.str_by_instant_cache[self] = instant_str = self.date.isoformat() + instant_str = t.InstantStr(self.date.isoformat()) + config.str_by_instant_cache[self] = instant_str return instant_str @property - def date(self): + def date(self) -> Date: instant_date = config.date_by_instant_cache.get(self) if instant_date is None: - config.date_by_instant_cache[self] = instant_date = pendulum.date(*self) + instant_date = pendulum.date(*self) + config.date_by_instant_cache[self] = instant_date return instant_date @property - def day(self): + def day(self) -> int: return self[2] @property - def month(self): + def month(self) -> int: return self[1] - def offset(self, offset, unit): + @property + def year(self) -> int: + return self[0] + + def offset(self, offset: str | int, unit: t.DateUnit) -> t.Instant | None: """Increments/decrements the given instant with offset units. Args: @@ -193,6 +199,9 @@ def offset(self, offset, unit): date = date.add(days=offset) return self.__class__((date.year, date.month, date.day)) - @property - def year(self): - return self[0] + return None + + @classmethod + def eternity(cls) -> t.Instant: + """Return an eternity instant.""" + return cls((1, 1, 1)) diff --git a/openfisca_core/periods/period_.py b/openfisca_core/periods/period_.py index 6dfb25873..5dee6f189 100644 --- a/openfisca_core/periods/period_.py +++ b/openfisca_core/periods/period_.py @@ -1,22 +1,19 @@ from __future__ import annotations -import typing from collections.abc import Sequence import calendar import datetime -import pendulum +from pendulum.datetime import Date from . import helpers +from . import types as t from .date_unit import DateUnit from .instant_ import Instant -if typing.TYPE_CHECKING: - from pendulum.datetime import Date - -class Period(tuple): +class Period(tuple[t.DateUnit, t.Instant, int]): """Toolbox to handle date intervals. A :class:`.Period` is a triple (``unit``, ``start``, ``size``). @@ -124,11 +121,11 @@ class Period(tuple): def __repr__(self) -> str: return f"{self.__class__.__name__}({super().__repr__()})" - def __str__(self) -> str: + def __str__(self) -> t.PeriodStr: unit, start_instant, size = self if unit == DateUnit.ETERNITY: - return unit.upper() + return t.PeriodStr(unit.upper()) # ISO format date units. f_year, month, day = start_instant @@ -140,58 +137,58 @@ def __str__(self) -> str: if unit == DateUnit.MONTH and size == 12 or unit == DateUnit.YEAR and size == 1: if month == 1: # civil year starting from january - return str(f_year) + return t.PeriodStr(str(f_year)) else: # rolling year - return f"{DateUnit.YEAR}:{f_year}-{month:02d}" + return t.PeriodStr(f"{DateUnit.YEAR}:{f_year}-{month:02d}") # simple month if unit == DateUnit.MONTH and size == 1: - return f"{f_year}-{month:02d}" + return t.PeriodStr(f"{f_year}-{month:02d}") # several civil years if unit == DateUnit.YEAR and month == 1: - return f"{unit}:{f_year}:{size}" + return t.PeriodStr(f"{unit}:{f_year}:{size}") if unit == DateUnit.DAY: if size == 1: - return f"{f_year}-{month:02d}-{day:02d}" + return t.PeriodStr(f"{f_year}-{month:02d}-{day:02d}") else: - return f"{unit}:{f_year}-{month:02d}-{day:02d}:{size}" + return t.PeriodStr(f"{unit}:{f_year}-{month:02d}-{day:02d}:{size}") # 1 week if unit == DateUnit.WEEK and size == 1: if week < 10: - return f"{c_year}-W0{week}" + return t.PeriodStr(f"{c_year}-W0{week}") - return f"{c_year}-W{week}" + return t.PeriodStr(f"{c_year}-W{week}") # several weeks if unit == DateUnit.WEEK and size > 1: if week < 10: - return f"{unit}:{c_year}-W0{week}:{size}" + return t.PeriodStr(f"{unit}:{c_year}-W0{week}:{size}") - return f"{unit}:{c_year}-W{week}:{size}" + return t.PeriodStr(f"{unit}:{c_year}-W{week}:{size}") # 1 weekday if unit == DateUnit.WEEKDAY and size == 1: if week < 10: - return f"{c_year}-W0{week}-{weekday}" + return t.PeriodStr(f"{c_year}-W0{week}-{weekday}") - return f"{c_year}-W{week}-{weekday}" + return t.PeriodStr(f"{c_year}-W{week}-{weekday}") # several weekdays if unit == DateUnit.WEEKDAY and size > 1: if week < 10: - return f"{unit}:{c_year}-W0{week}-{weekday}:{size}" + return t.PeriodStr(f"{unit}:{c_year}-W0{week}-{weekday}:{size}") - return f"{unit}:{c_year}-W{week}-{weekday}:{size}" + return t.PeriodStr(f"{unit}:{c_year}-W{week}-{weekday}:{size}") # complex period - return f"{unit}:{f_year}-{month:02d}:{size}" + return t.PeriodStr(f"{unit}:{f_year}-{month:02d}:{size}") @property - def unit(self) -> str: + def unit(self) -> t.DateUnit: """The ``unit`` of the ``Period``. Example: @@ -205,7 +202,7 @@ def unit(self) -> str: return self[0] @property - def start(self) -> Instant: + def start(self) -> t.Instant: """The ``Instant`` at which the ``Period`` starts. Example: @@ -334,7 +331,7 @@ def size_in_days(self) -> int: raise ValueError(f"Can't calculate number of days in a {self.unit}.") @property - def size_in_weeks(self): + def size_in_weeks(self) -> int: """The ``size`` of the ``Period`` in weeks. Examples: @@ -353,13 +350,13 @@ def size_in_weeks(self): if self.unit == DateUnit.YEAR: start = self.start.date cease = start.add(years=self.size) - delta = pendulum.interval(start, cease) + delta = start.diff(cease) return delta.in_weeks() if self.unit == DateUnit.MONTH: start = self.start.date cease = start.add(months=self.size) - delta = pendulum.interval(start, cease) + delta = start.diff(cease) return delta.in_weeks() if self.unit == DateUnit.WEEK: @@ -368,7 +365,7 @@ def size_in_weeks(self): raise ValueError(f"Can't calculate number of weeks in a {self.unit}.") @property - def size_in_weekdays(self): + def size_in_weekdays(self) -> int: """The ``size`` of the ``Period`` in weekdays. Examples: @@ -400,11 +397,13 @@ def size_in_weekdays(self): raise ValueError(f"Can't calculate number of weekdays in a {self.unit}.") @property - def days(self): + def days(self) -> int: """Same as ``size_in_days``.""" return (self.stop.date - self.start.date).days + 1 - def intersection(self, start, stop): + def intersection( + self, start: t.Instant | None, stop: t.Instant | None + ) -> t.Period | None: if start is None and stop is None: return self period_start = self[1] @@ -457,7 +456,7 @@ def intersection(self, start, stop): ) ) - def get_subperiods(self, unit: DateUnit) -> Sequence[Period]: + def get_subperiods(self, unit: t.DateUnit) -> Sequence[t.Period]: """Return the list of periods of unit ``unit`` contained in self. Examples: @@ -502,7 +501,7 @@ def get_subperiods(self, unit: DateUnit) -> Sequence[Period]: raise ValueError(f"Cannot subdivide {self.unit} into {unit}") - def offset(self, offset, unit=None): + def offset(self, offset: str | int, unit: t.DateUnit | None = None) -> t.Period: """Increment (or decrement) the given period with offset units. Examples: @@ -684,7 +683,7 @@ def offset(self, offset, unit=None): ) ) - def contains(self, other: Period) -> bool: + def contains(self, other: t.Period) -> bool: """Returns ``True`` if the period contains ``other``. For instance, ``period(2015)`` contains ``period(2015-01)``. @@ -694,7 +693,7 @@ def contains(self, other: Period) -> bool: return self.start <= other.start and self.stop >= other.stop @property - def stop(self) -> Instant: + def stop(self) -> t.Instant: """Return the last day of the period as an Instant instance. Examples: @@ -731,7 +730,7 @@ def stop(self) -> Instant: year, month, day = start_instant if unit == DateUnit.ETERNITY: - return Instant((float("inf"), float("inf"), float("inf"))) + return Instant.eternity() elif unit == DateUnit.YEAR: date = start_instant.date.add(years=size, days=-1) @@ -755,67 +754,72 @@ def stop(self) -> Instant: # Reference periods @property - def last_week(self) -> Period: + def last_week(self) -> t.Period: return self.first_week.offset(-1) @property - def last_fortnight(self) -> Period: - start: Instant = self.first_week.start + def last_fortnight(self) -> t.Period: + start: t.Instant = self.first_week.start return self.__class__((DateUnit.WEEK, start, 1)).offset(-2) @property - def last_2_weeks(self) -> Period: - start: Instant = self.first_week.start + def last_2_weeks(self) -> t.Period: + start: t.Instant = self.first_week.start return self.__class__((DateUnit.WEEK, start, 2)).offset(-2) @property - def last_26_weeks(self) -> Period: - start: Instant = self.first_week.start + def last_26_weeks(self) -> t.Period: + start: t.Instant = self.first_week.start return self.__class__((DateUnit.WEEK, start, 26)).offset(-26) @property - def last_52_weeks(self) -> Period: - start: Instant = self.first_week.start + def last_52_weeks(self) -> t.Period: + start: t.Instant = self.first_week.start return self.__class__((DateUnit.WEEK, start, 52)).offset(-52) @property - def last_month(self) -> Period: + def last_month(self) -> t.Period: return self.first_month.offset(-1) @property - def last_3_months(self) -> Period: - start: Instant = self.first_month.start + def last_3_months(self) -> t.Period: + start: t.Instant = self.first_month.start return self.__class__((DateUnit.MONTH, start, 3)).offset(-3) @property - def last_year(self) -> Period: - start: Instant = self.start.offset("first-of", DateUnit.YEAR) + def last_year(self) -> t.Period: + start: t.Instant = self.start.offset("first-of", DateUnit.YEAR) return self.__class__((DateUnit.YEAR, start, 1)).offset(-1) @property - def n_2(self) -> Period: - start: Instant = self.start.offset("first-of", DateUnit.YEAR) + def n_2(self) -> t.Period: + start: t.Instant = self.start.offset("first-of", DateUnit.YEAR) return self.__class__((DateUnit.YEAR, start, 1)).offset(-2) @property - def this_year(self) -> Period: - start: Instant = self.start.offset("first-of", DateUnit.YEAR) + def this_year(self) -> t.Period: + start: t.Instant = self.start.offset("first-of", DateUnit.YEAR) return self.__class__((DateUnit.YEAR, start, 1)) @property - def first_month(self) -> Period: - start: Instant = self.start.offset("first-of", DateUnit.MONTH) + def first_month(self) -> t.Period: + start: t.Instant = self.start.offset("first-of", DateUnit.MONTH) return self.__class__((DateUnit.MONTH, start, 1)) @property - def first_day(self) -> Period: + def first_day(self) -> t.Period: return self.__class__((DateUnit.DAY, self.start, 1)) @property - def first_week(self) -> Period: - start: Instant = self.start.offset("first-of", DateUnit.WEEK) + def first_week(self) -> t.Period: + start: t.Instant = self.start.offset("first-of", DateUnit.WEEK) return self.__class__((DateUnit.WEEK, start, 1)) @property - def first_weekday(self) -> Period: + def first_weekday(self) -> t.Period: return self.__class__((DateUnit.WEEKDAY, self.start, 1)) + + @classmethod + def eternity(cls) -> t.Period: + """Return an eternity period.""" + return cls((DateUnit.ETERNITY, Instant.eternity(), 0)) diff --git a/openfisca_core/periods/tests/helpers/test_helpers.py b/openfisca_core/periods/tests/helpers/test_helpers.py index bb409323d..1a351ce83 100644 --- a/openfisca_core/periods/tests/helpers/test_helpers.py +++ b/openfisca_core/periods/tests/helpers/test_helpers.py @@ -47,9 +47,19 @@ def test_instant_date_with_an_invalid_argument(arg, error): [Period((DateUnit.MONTH, Instant((1, 1, 1)), 12)), "200_12"], [Period((DateUnit.YEAR, Instant((1, 1, 1)), 2)), "300_2"], [Period((DateUnit.ETERNITY, Instant((1, 1, 1)), 1)), "400_1"], - [(DateUnit.DAY, None, 1), "100_1"], - [(DateUnit.MONTH, None, -1000), "200_-1000"], ], ) def test_key_period_size(arg, expected): assert periods.key_period_size(arg) == expected + + +@pytest.mark.parametrize( + "arg, error", + [ + [(DateUnit.DAY, None, 1), AttributeError], + [(DateUnit.MONTH, None, -1000), AttributeError], + ], +) +def test_key_period_size_when_an_invalid_argument(arg, error): + with pytest.raises(error): + periods.key_period_size(arg) diff --git a/openfisca_core/periods/tests/helpers/test_period.py b/openfisca_core/periods/tests/helpers/test_period.py index 7d50abe10..33aaac2e2 100644 --- a/openfisca_core/periods/tests/helpers/test_period.py +++ b/openfisca_core/periods/tests/helpers/test_period.py @@ -9,11 +9,11 @@ @pytest.mark.parametrize( "arg, expected", [ - ["eternity", Period((DateUnit.ETERNITY, Instant((1, 1, 1)), float("inf")))], - ["ETERNITY", Period((DateUnit.ETERNITY, Instant((1, 1, 1)), float("inf")))], + ["eternity", Period((DateUnit.ETERNITY, Instant((1, 1, 1)), 0))], + ["ETERNITY", Period((DateUnit.ETERNITY, Instant((1, 1, 1)), 0))], [ DateUnit.ETERNITY, - Period((DateUnit.ETERNITY, Instant((1, 1, 1)), float("inf"))), + Period((DateUnit.ETERNITY, Instant((1, 1, 1)), 0)), ], [datetime.date(1, 1, 1), Period((DateUnit.DAY, Instant((1, 1, 1)), 1))], [Instant((1, 1, 1)), Period((DateUnit.DAY, Instant((1, 1, 1)), 1))], diff --git a/openfisca_core/periods/types.py b/openfisca_core/periods/types.py index e69de29bb..f8136f5b6 100644 --- a/openfisca_core/periods/types.py +++ b/openfisca_core/periods/types.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from typing import NewType, Protocol + +from pendulum.datetime import Date + +from openfisca_core import types as t + +# New types. + +#: For example "2000-01". +InstantStr = NewType("InstantStr", str) + +#: For example "1:2000-01-01:day". +PeriodStr = NewType("PeriodStr", str) + + +# Periods + + +class DateUnit(t.DateUnit, Protocol): + ... + + +class Instant(t.Instant, Protocol): + @property + def year(self) -> int: + ... + + @property + def month(self) -> int: + ... + + @property + def day(self) -> int: + ... + + @property + def date(self) -> Date: + ... + + def offset(self, offset: str | int, unit: DateUnit) -> Instant | None: + ... + + +class Period(t.Period, Protocol): + @property + def size(self) -> int: + ... + + @property + def start(self) -> Instant: + ... + + @property + def stop(self) -> Instant: + ... + + def offset(self, offset: str | int, unit: DateUnit | None = None) -> Period: + ... From 1392bb8672901c2aec5752b18f6828a649aaf16d Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 13:08:51 +0200 Subject: [PATCH 06/20] fix(periods): unhandled none --- openfisca_core/periods/period_.py | 34 ++++++++++++++++++++++++------- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/openfisca_core/periods/period_.py b/openfisca_core/periods/period_.py index 5dee6f189..700652953 100644 --- a/openfisca_core/periods/period_.py +++ b/openfisca_core/periods/period_.py @@ -319,7 +319,12 @@ def size_in_days(self) -> int: """ if self.unit in (DateUnit.YEAR, DateUnit.MONTH): - last_day = self.start.offset(self.size, self.unit).offset(-1, DateUnit.DAY) + last = self.start.offset(self.size, self.unit) + if last is None: + raise NotImplementedError + last_day = last.offset(-1, DateUnit.DAY) + if last_day is None: + raise NotImplementedError return (last_day.date - self.start.date).days + 1 if self.unit == DateUnit.WEEK: @@ -385,7 +390,12 @@ def size_in_weekdays(self) -> int: return self.size_in_weeks * 7 if self.unit in DateUnit.MONTH: - last_day = self.start.offset(self.size, self.unit).offset(-1, DateUnit.DAY) + last = self.start.offset(self.size, self.unit) + if last is None: + raise NotImplementedError + last_day = last.offset(-1, DateUnit.DAY) + if last_day is None: + raise NotImplementedError return (last_day.date - self.start.date).days + 1 if self.unit == DateUnit.WEEK: @@ -788,22 +798,30 @@ def last_3_months(self) -> t.Period: @property def last_year(self) -> t.Period: - start: t.Instant = self.start.offset("first-of", DateUnit.YEAR) + start: None | t.Instant = self.start.offset("first-of", DateUnit.YEAR) + if start is None: + raise NotImplementedError return self.__class__((DateUnit.YEAR, start, 1)).offset(-1) @property def n_2(self) -> t.Period: - start: t.Instant = self.start.offset("first-of", DateUnit.YEAR) + start: None | t.Instant = self.start.offset("first-of", DateUnit.YEAR) + if start is None: + raise NotImplementedError return self.__class__((DateUnit.YEAR, start, 1)).offset(-2) @property def this_year(self) -> t.Period: - start: t.Instant = self.start.offset("first-of", DateUnit.YEAR) + start: None | t.Instant = self.start.offset("first-of", DateUnit.YEAR) + if start is None: + raise NotImplementedError return self.__class__((DateUnit.YEAR, start, 1)) @property def first_month(self) -> t.Period: - start: t.Instant = self.start.offset("first-of", DateUnit.MONTH) + start: None | t.Instant = self.start.offset("first-of", DateUnit.MONTH) + if start is None: + raise NotImplementedError return self.__class__((DateUnit.MONTH, start, 1)) @property @@ -812,7 +830,9 @@ def first_day(self) -> t.Period: @property def first_week(self) -> t.Period: - start: t.Instant = self.start.offset("first-of", DateUnit.WEEK) + start: None | t.Instant = self.start.offset("first-of", DateUnit.WEEK) + if start is None: + raise NotImplementedError return self.__class__((DateUnit.WEEK, start, 1)) @property From e4326f6ac05239b3a074f06d9ce0cb9651f0637a Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 13:53:59 +0200 Subject: [PATCH 07/20] fix(periods): dunder method types --- openfisca_core/periods/date_unit.py | 5 +++++ openfisca_core/periods/instant_.py | 10 ++++++++++ openfisca_core/periods/period_.py | 11 +++++++++-- openfisca_core/periods/types.py | 6 ++++++ openfisca_core/types.py | 10 +++------- 5 files changed, 33 insertions(+), 9 deletions(-) diff --git a/openfisca_core/periods/date_unit.py b/openfisca_core/periods/date_unit.py index a808b1941..13d93f8a2 100644 --- a/openfisca_core/periods/date_unit.py +++ b/openfisca_core/periods/date_unit.py @@ -98,6 +98,11 @@ class DateUnit(StrEnum, metaclass=DateUnitMeta): """ + def __contains__(self, other: object) -> bool: + if isinstance(other, str): + return super().__contains__(other) + return NotImplemented + WEEKDAY = "weekday" WEEK = "week" DAY = "day" diff --git a/openfisca_core/periods/instant_.py b/openfisca_core/periods/instant_.py index e68151c36..5eed50c45 100644 --- a/openfisca_core/periods/instant_.py +++ b/openfisca_core/periods/instant_.py @@ -90,6 +90,16 @@ def __str__(self) -> t.InstantStr: return instant_str + def __lt__(self, other: object) -> bool: + if isinstance(other, Instant): + return super().__lt__(other) + return NotImplemented + + def __le__(self, other: object) -> bool: + if isinstance(other, Instant): + return super().__le__(other) + return NotImplemented + @property def date(self) -> Date: instant_date = config.date_by_instant_cache.get(self) diff --git a/openfisca_core/periods/period_.py b/openfisca_core/periods/period_.py index 700652953..39884d900 100644 --- a/openfisca_core/periods/period_.py +++ b/openfisca_core/periods/period_.py @@ -389,7 +389,7 @@ def size_in_weekdays(self) -> int: if self.unit == DateUnit.YEAR: return self.size_in_weeks * 7 - if self.unit in DateUnit.MONTH: + if DateUnit.MONTH in self.unit: last = self.start.offset(self.size, self.unit) if last is None: raise NotImplementedError @@ -685,10 +685,17 @@ def offset(self, offset: str | int, unit: t.DateUnit | None = None) -> t.Period: """ + start: None | t.Instant = self[1].offset( + offset, self[0] if unit is None else unit + ) + + if start is None: + raise NotImplementedError + return self.__class__( ( self[0], - self[1].offset(offset, self[0] if unit is None else unit), + start, self[2], ) ) diff --git a/openfisca_core/periods/types.py b/openfisca_core/periods/types.py index f8136f5b6..26a6e096d 100644 --- a/openfisca_core/periods/types.py +++ b/openfisca_core/periods/types.py @@ -39,6 +39,12 @@ def day(self) -> int: def date(self) -> Date: ... + def __lt__(self, other: object, /) -> bool: + ... + + def __le__(self, other: object, /) -> bool: + ... + def offset(self, offset: str | int, unit: DateUnit) -> Instant | None: ... diff --git a/openfisca_core/types.py b/openfisca_core/types.py index fe9ea8dee..2c1e9945e 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -88,18 +88,14 @@ class ParameterNodeAtInstant(Protocol): # Periods -class Container(Protocol[T_con]): - def __contains__(self, item: T_con, /) -> bool: - ... - - class Indexable(Protocol[T_cov]): def __getitem__(self, index: int, /) -> T_cov: ... -class DateUnit(Container[str], Protocol): - ... +class DateUnit(Protocol): + def __contains__(self, other: object, /) -> bool: + ... class Instant(Indexable[int], Iterable[int], Sized, Protocol): From 101b9a18e479af00367c878db31f32e33e79b20d Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 13:54:44 +0200 Subject: [PATCH 08/20] chore(periods): add py.typed --- openfisca_core/periods/py.typed | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 openfisca_core/periods/py.typed diff --git a/openfisca_core/periods/py.typed b/openfisca_core/periods/py.typed new file mode 100644 index 000000000..e69de29bb From eb159e80c71cb3f47206d6a641a418e84278ad18 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 00:34:58 +0200 Subject: [PATCH 09/20] chore(version): bump --- CHANGELOG.md | 8 ++++++++ setup.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ade861023..8d3f5ff6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +### 41.5.6 [#1223](https://github.com/openfisca/openfisca-core/pull/1223) + +#### Technical changes + +- Update `pendulum' +- Remove run-time type-checks +- Add typing to the periods module + ### 41.5.5 [#1220](https://github.com/openfisca/openfisca-core/pull/1220) #### Technical changes diff --git a/setup.py b/setup.py index 6707ce9c6..d62443e9e 100644 --- a/setup.py +++ b/setup.py @@ -67,7 +67,7 @@ setup( name="OpenFisca-Core", - version="41.5.5", + version="41.5.6", author="OpenFisca Team", author_email="contact@openfisca.org", classifiers=[ From a919dfd30a7592debba36ddd865143e69628b855 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 14:50:23 +0200 Subject: [PATCH 10/20] feat(enums): add specific types for numpy arrays --- openfisca_core/indexed_enums/types.py | 14 ++++++++++ openfisca_core/types.py | 39 +++++++++++++++++++++++++-- 2 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 openfisca_core/indexed_enums/types.py diff --git a/openfisca_core/indexed_enums/types.py b/openfisca_core/indexed_enums/types.py new file mode 100644 index 000000000..592c9fab7 --- /dev/null +++ b/openfisca_core/indexed_enums/types.py @@ -0,0 +1,14 @@ +from typing import Protocol + +from openfisca_core.types import Array, ArrayBytes, ArrayEnum, ArrayInt, ArrayStr + +# Indexed enums + + +class Enum(Protocol): + @classmethod + def encode(cls, array: object) -> object: + ... + + +__all__ = ["Array", "ArrayEnum", "ArrayStr", "ArrayBytes", "ArrayInt"] diff --git a/openfisca_core/types.py b/openfisca_core/types.py index 2c1e9945e..44780fb7b 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -2,16 +2,49 @@ from collections.abc import Iterable, Sequence, Sized from numpy.typing import NDArray -from typing import Any, TypeVar, Union +from typing import Any, NewType, TypeVar, Union from typing_extensions import Protocol, TypeAlias import numpy +# New types + +#: Type for arrays of any type. +ArrayAny = NewType("ArrayAny", numpy.generic) + +#: Type for arrays of booleans. +ArrayBool = NewType("ArrayBool", numpy.bool_) + +#: Type for arrays of bytes. +ArrayBytes = NewType("ArrayBytes", numpy.bytes_) + +#: Type for arrays of dates. +ArrayDate = NewType("ArrayDate", numpy.datetime64) + +#: Type for arrays of enums. +ArrayEnum = NewType("ArrayEnum", numpy.int16) + +#: Type for arrays of floats. +ArrayFloat = NewType("ArrayFloat", numpy.float32) + +#: Type for arrays of integers. +ArrayInt = NewType("ArrayInt", numpy.int32) + +#: Type for arrays of Python objects. +ArrayObject = NewType("ArrayObject", numpy.object_) + +#: Type for arrays of strings. +ArrayStr = NewType("ArrayStr", numpy.str_) + +# TypeAliases + +#: Generic numpy type an array can have. N = TypeVar("N", bound=numpy.generic, covariant=True) #: Type representing an numpy array. Array: TypeAlias = NDArray[N] +#: Generic type a sequence can have. L = TypeVar("L") #: Type representing an array-like object. @@ -23,8 +56,10 @@ #: Type variable representing a value. A = TypeVar("A", covariant=True) -#: Generic type vars. +#: Generic covariant type. T_cov = TypeVar("T_cov", covariant=True) + +#: Generic contravariant type. T_con = TypeVar("T_con", contravariant=True) From 6646dd601d8f4a1a251a78605b8f627e58a65571 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 15:04:37 +0200 Subject: [PATCH 11/20] refactor(enums): reuse data types elsewhere --- openfisca_core/indexed_enums/__init__.py | 7 ++++--- openfisca_core/indexed_enums/config.py | 3 --- openfisca_core/indexed_enums/enum.py | 21 ++++++++++----------- openfisca_core/types.py | 24 +++++++++++------------- openfisca_core/variables/config.py | 16 +++++++--------- 5 files changed, 32 insertions(+), 39 deletions(-) delete mode 100644 openfisca_core/indexed_enums/config.py diff --git a/openfisca_core/indexed_enums/__init__.py b/openfisca_core/indexed_enums/__init__.py index 874a7e1f9..03958426b 100644 --- a/openfisca_core/indexed_enums/__init__.py +++ b/openfisca_core/indexed_enums/__init__.py @@ -21,6 +21,7 @@ # # See: https://www.python.org/dev/peps/pep-0008/#imports -from .config import ENUM_ARRAY_DTYPE # noqa: F401 -from .enum import Enum # noqa: F401 -from .enum_array import EnumArray # noqa: F401 +from .enum import Enum +from .enum_array import EnumArray + +__all__ = ["Enum", "EnumArray"] diff --git a/openfisca_core/indexed_enums/config.py b/openfisca_core/indexed_enums/config.py deleted file mode 100644 index f7da69b84..000000000 --- a/openfisca_core/indexed_enums/config.py +++ /dev/null @@ -1,3 +0,0 @@ -import numpy - -ENUM_ARRAY_DTYPE = numpy.int16 diff --git a/openfisca_core/indexed_enums/enum.py b/openfisca_core/indexed_enums/enum.py index 7957ced3a..a6fa37e6c 100644 --- a/openfisca_core/indexed_enums/enum.py +++ b/openfisca_core/indexed_enums/enum.py @@ -1,12 +1,10 @@ from __future__ import annotations -from typing import Union - import enum import numpy -from .config import ENUM_ARRAY_DTYPE +from . import types as t from .enum_array import EnumArray @@ -33,12 +31,13 @@ def __init__(self, name: str) -> None: @classmethod def encode( cls, - array: Union[ - EnumArray, - numpy.int_, - numpy.float_, - numpy.object_, - ], + array: ( + EnumArray + | t.Array[t.ArrayBytes] + | t.Array[t.ArrayEnum] + | t.Array[t.ArrayInt] + | t.Array[t.ArrayStr] + ), ) -> EnumArray: """ Encode a string numpy array, an enum item numpy array, or an int numpy @@ -73,7 +72,7 @@ def encode( array = numpy.select( [array == item.name for item in cls], [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) + ).astype(t.ArrayEnum) # Enum items arrays elif isinstance(array, numpy.ndarray) and array.dtype.kind == "O": @@ -94,6 +93,6 @@ def encode( array = numpy.select( [array == item for item in cls], [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) + ).astype(t.ArrayEnum) return EnumArray(array, cls) diff --git a/openfisca_core/types.py b/openfisca_core/types.py index 44780fb7b..772a0d29e 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -2,41 +2,39 @@ from collections.abc import Iterable, Sequence, Sized from numpy.typing import NDArray -from typing import Any, NewType, TypeVar, Union +from typing import Any, TypeVar, Union from typing_extensions import Protocol, TypeAlias import numpy -# New types +# Type aliases #: Type for arrays of any type. -ArrayAny = NewType("ArrayAny", numpy.generic) +ArrayAny: TypeAlias = numpy.generic #: Type for arrays of booleans. -ArrayBool = NewType("ArrayBool", numpy.bool_) +ArrayBool: TypeAlias = numpy.bool_ #: Type for arrays of bytes. -ArrayBytes = NewType("ArrayBytes", numpy.bytes_) +ArrayBytes: TypeAlias = numpy.bytes_ #: Type for arrays of dates. -ArrayDate = NewType("ArrayDate", numpy.datetime64) +ArrayDate: TypeAlias = numpy.datetime64 #: Type for arrays of enums. -ArrayEnum = NewType("ArrayEnum", numpy.int16) +ArrayEnum: TypeAlias = numpy.int16 #: Type for arrays of floats. -ArrayFloat = NewType("ArrayFloat", numpy.float32) +ArrayFloat: TypeAlias = numpy.float32 #: Type for arrays of integers. -ArrayInt = NewType("ArrayInt", numpy.int32) +ArrayInt: TypeAlias = numpy.int32 #: Type for arrays of Python objects. -ArrayObject = NewType("ArrayObject", numpy.object_) +ArrayObject: TypeAlias = numpy.object_ #: Type for arrays of strings. -ArrayStr = NewType("ArrayStr", numpy.str_) - -# TypeAliases +ArrayStr: TypeAlias = numpy.str_ #: Generic numpy type an array can have. N = TypeVar("N", bound=numpy.generic, covariant=True) diff --git a/openfisca_core/variables/config.py b/openfisca_core/variables/config.py index 54270145b..90faa7d12 100644 --- a/openfisca_core/variables/config.py +++ b/openfisca_core/variables/config.py @@ -1,27 +1,25 @@ import datetime -import numpy - -from openfisca_core import indexed_enums -from openfisca_core.indexed_enums import Enum +from openfisca_core import indexed_enums as enums +from openfisca_core import types as t VALUE_TYPES = { bool: { - "dtype": numpy.bool_, + "dtype": t.ArrayBool, "default": False, "json_type": "boolean", "formatted_value_type": "Boolean", "is_period_size_independent": True, }, int: { - "dtype": numpy.int32, + "dtype": t.ArrayInt, "default": 0, "json_type": "integer", "formatted_value_type": "Int", "is_period_size_independent": False, }, float: { - "dtype": numpy.float32, + "dtype": t.ArrayFloat, "default": 0, "json_type": "number", "formatted_value_type": "Float", @@ -34,8 +32,8 @@ "formatted_value_type": "String", "is_period_size_independent": True, }, - Enum: { - "dtype": indexed_enums.ENUM_ARRAY_DTYPE, + enums.Enum: { + "dtype": t.ArrayEnum, "json_type": "string", "formatted_value_type": "String", "is_period_size_independent": True, From 63a5c28c73fc7af32f58d868ab01bbd2e25ebbb0 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 17:00:19 +0200 Subject: [PATCH 12/20] doc(enums): add types to __array_finalize__ --- openfisca_core/indexed_enums/enum.py | 4 ++-- openfisca_core/indexed_enums/enum_array.py | 28 +++++++++++++--------- openfisca_core/indexed_enums/types.py | 24 +++++++++++++++++-- 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/openfisca_core/indexed_enums/enum.py b/openfisca_core/indexed_enums/enum.py index a6fa37e6c..1b09a93c2 100644 --- a/openfisca_core/indexed_enums/enum.py +++ b/openfisca_core/indexed_enums/enum.py @@ -32,13 +32,13 @@ def __init__(self, name: str) -> None: def encode( cls, array: ( - EnumArray + t.EnumArray | t.Array[t.ArrayBytes] | t.Array[t.ArrayEnum] | t.Array[t.ArrayInt] | t.Array[t.ArrayStr] ), - ) -> EnumArray: + ) -> t.EnumArray: """ Encode a string numpy array, an enum item numpy array, or an int numpy array into an :any:`EnumArray`. See :any:`EnumArray.decode` for diff --git a/openfisca_core/indexed_enums/enum_array.py b/openfisca_core/indexed_enums/enum_array.py index 2742719ad..7d2e7e711 100644 --- a/openfisca_core/indexed_enums/enum_array.py +++ b/openfisca_core/indexed_enums/enum_array.py @@ -1,15 +1,13 @@ from __future__ import annotations -import typing -from typing import Any, NoReturn, Optional, Type +from typing import Any, NoReturn import numpy -if typing.TYPE_CHECKING: - from openfisca_core.indexed_enums import Enum +from . import types as t -class EnumArray(numpy.ndarray): +class EnumArray(t.Array[t.ArrayEnum]): """ NumPy array subclass representing an array of enum items. @@ -19,23 +17,31 @@ class EnumArray(numpy.ndarray): # Subclassing ndarray is a little tricky. # To read more about the two following methods, see: # https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array. + + #: Enum type of the array items. + possible_values: type[t.Enum] | None = None + def __new__( cls, - input_array: numpy.int_, - possible_values: Optional[Type[Enum]] = None, + input_array: t.Array[t.ArrayEnum], + possible_values: type[t.Enum] | None = None, ) -> EnumArray: obj = numpy.asarray(input_array).view(cls) obj.possible_values = possible_values return obj # See previous comment - def __array_finalize__(self, obj: Optional[numpy.int_]) -> None: + def __array_finalize__(self, obj: t.Array[t.ArrayAny] | t.EnumArray | None) -> None: if obj is None: - return + return None + if isinstance(obj, EnumArray): + self.possible_values = obj.possible_values + return None - self.possible_values = getattr(obj, "possible_values", None) + def __eq__(self, other: object) -> t.Array[t.ArrayBool] | bool: # type: ignore[override] + if self.possible_values is None: + raise NotImplementedError - def __eq__(self, other: Any) -> bool: # When comparing to an item of self.possible_values, use the item index # to speed up the comparison. if other.__class__.__name__ is self.possible_values.__name__: diff --git a/openfisca_core/indexed_enums/types.py b/openfisca_core/indexed_enums/types.py index 592c9fab7..0aa060b85 100644 --- a/openfisca_core/indexed_enums/types.py +++ b/openfisca_core/indexed_enums/types.py @@ -1,6 +1,14 @@ from typing import Protocol -from openfisca_core.types import Array, ArrayBytes, ArrayEnum, ArrayInt, ArrayStr +from openfisca_core.types import ( + Array, + ArrayAny, + ArrayBool, + ArrayBytes, + ArrayEnum, + ArrayInt, + ArrayStr, +) # Indexed enums @@ -11,4 +19,16 @@ def encode(cls, array: object) -> object: ... -__all__ = ["Array", "ArrayEnum", "ArrayStr", "ArrayBytes", "ArrayInt"] +class EnumArray(Protocol): + ... + + +__all__ = [ + "Array", + "ArrayAny", + "ArrayBool", + "ArrayEnum", + "ArrayStr", + "ArrayBytes", + "ArrayInt", +] From 96800b8e04b7948bddd4b577e4579755acbafc25 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 01:42:08 +0200 Subject: [PATCH 13/20] fix(enums): handle complex type override --- openfisca_core/indexed_enums/enum_array.py | 84 ++++++++++++++++------ openfisca_core/indexed_enums/types.py | 35 ++++----- 2 files changed, 76 insertions(+), 43 deletions(-) diff --git a/openfisca_core/indexed_enums/enum_array.py b/openfisca_core/indexed_enums/enum_array.py index 7d2e7e711..7563ab251 100644 --- a/openfisca_core/indexed_enums/enum_array.py +++ b/openfisca_core/indexed_enums/enum_array.py @@ -1,13 +1,14 @@ from __future__ import annotations -from typing import Any, NoReturn +from typing import NoReturn, overload +from typing_extensions import TypeGuard import numpy from . import types as t -class EnumArray(t.Array[t.ArrayEnum]): +class EnumArray(t.EnumArray): """ NumPy array subclass representing an array of enum items. @@ -19,57 +20,84 @@ class EnumArray(t.Array[t.ArrayEnum]): # https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array. #: Enum type of the array items. - possible_values: type[t.Enum] | None = None + possible_values: None | type[t.Enum] = None def __new__( cls, input_array: t.Array[t.ArrayEnum], - possible_values: type[t.Enum] | None = None, + possible_values: None | type[t.Enum] = None, ) -> EnumArray: obj = numpy.asarray(input_array).view(cls) obj.possible_values = possible_values return obj # See previous comment - def __array_finalize__(self, obj: t.Array[t.ArrayAny] | t.EnumArray | None) -> None: + def __array_finalize__(self, obj: None | t.EnumArray | t.Array[t.ArrayAny]) -> None: if obj is None: return None if isinstance(obj, EnumArray): self.possible_values = obj.possible_values return None - def __eq__(self, other: object) -> t.Array[t.ArrayBool] | bool: # type: ignore[override] + @overload # type: ignore[override] + def __eq__(self, other: None | t.Enum | type[t.Enum]) -> t.Array[t.ArrayBool]: + ... + + @overload + def __eq__(self, other: object) -> t.Array[t.ArrayBool] | bool: + ... + + def __eq__(self, other: object) -> t.Array[t.ArrayBool] | bool: + boolean_array: t.Array[t.ArrayBool] + boolean: bool + if self.possible_values is None: - raise NotImplementedError + return NotImplemented + + view: t.Array[t.ArrayEnum] = self.view(numpy.ndarray) + + if other is None or self._is_an_enum_type(other): + boolean_array = view == other + return boolean_array # When comparing to an item of self.possible_values, use the item index # to speed up the comparison. - if other.__class__.__name__ is self.possible_values.__name__: + if self._is_an_enum(other): # Use view(ndarray) so that the result is a classic ndarray, not an # EnumArray. - return self.view(numpy.ndarray) == other.index + boolean_array = view == other.index + return boolean_array + + boolean = view == other + return boolean - return self.view(numpy.ndarray) == other + @overload # type: ignore[override] + def __ne__(self, other: None | t.Enum | type[t.Enum]) -> t.Array[t.ArrayBool]: + ... - def __ne__(self, other: Any) -> bool: + @overload + def __ne__(self, other: object) -> t.Array[t.ArrayBool] | bool: + ... + + def __ne__(self, other: object) -> t.Array[t.ArrayBool] | bool: return numpy.logical_not(self == other) - def _forbidden_operation(self, other: Any) -> NoReturn: + def _forbidden_operation(self, other: object) -> NoReturn: raise TypeError( "Forbidden operation. The only operations allowed on EnumArrays " "are '==' and '!='.", ) - __add__ = _forbidden_operation - __mul__ = _forbidden_operation - __lt__ = _forbidden_operation - __le__ = _forbidden_operation - __gt__ = _forbidden_operation - __ge__ = _forbidden_operation - __and__ = _forbidden_operation - __or__ = _forbidden_operation + __add__ = _forbidden_operation # type: ignore[assignment] + __mul__ = _forbidden_operation # type: ignore[assignment] + __lt__ = _forbidden_operation # type: ignore[assignment] + __le__ = _forbidden_operation # type: ignore[assignment] + __gt__ = _forbidden_operation # type: ignore[assignment] + __ge__ = _forbidden_operation # type: ignore[assignment] + __and__ = _forbidden_operation # type: ignore[assignment] + __or__ = _forbidden_operation # type: ignore[assignment] - def decode(self) -> numpy.object_: + def decode(self) -> t.Array[t.ArrayEnum]: """ Return the array of enum items corresponding to self. @@ -88,7 +116,7 @@ def decode(self) -> numpy.object_: list(self.possible_values), ) - def decode_to_str(self) -> numpy.str_: + def decode_to_str(self) -> t.Array[t.ArrayStr]: """ Return the array of string identifiers corresponding to self. @@ -110,3 +138,15 @@ def __repr__(self) -> str: def __str__(self) -> str: return str(self.decode_to_str()) + + def _is_an_enum(self, other: object) -> TypeGuard[t.Enum]: + return ( + not hasattr(other, "__name__") + and other.__class__.__name__ is self.possible_values.__name__ + ) + + def _is_an_enum_type(self, other: object) -> TypeGuard[type[t.Enum]]: + return ( + hasattr(other, "__name__") + and other.__name__ is self.possible_values.__name__ + ) diff --git a/openfisca_core/indexed_enums/types.py b/openfisca_core/indexed_enums/types.py index 0aa060b85..f8670af0f 100644 --- a/openfisca_core/indexed_enums/types.py +++ b/openfisca_core/indexed_enums/types.py @@ -1,34 +1,27 @@ +from __future__ import annotations + from typing import Protocol -from openfisca_core.types import ( - Array, - ArrayAny, - ArrayBool, - ArrayBytes, - ArrayEnum, - ArrayInt, - ArrayStr, -) +from openfisca_core.types import Array as Array +from openfisca_core.types import ArrayAny as ArrayAny # noqa: F401 +from openfisca_core.types import ArrayBool as ArrayBool # noqa: F401 +from openfisca_core.types import ArrayBytes as ArrayBytes # noqa: F401 +from openfisca_core.types import ArrayEnum as ArrayEnum +from openfisca_core.types import ArrayInt as ArrayInt # noqa: F401 +from openfisca_core.types import ArrayStr as ArrayStr # noqa: F401 + +import abc # Indexed enums class Enum(Protocol): + index: int + @classmethod def encode(cls, array: object) -> object: ... -class EnumArray(Protocol): +class EnumArray(Array[ArrayEnum], metaclass=abc.ABCMeta): ... - - -__all__ = [ - "Array", - "ArrayAny", - "ArrayBool", - "ArrayEnum", - "ArrayStr", - "ArrayBytes", - "ArrayInt", -] From 17bfb42530956dfca60336606985dbc810941ad9 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 03:39:55 +0200 Subject: [PATCH 14/20] fix(enums): broken types --- openfisca_core/indexed_enums/enum.py | 8 ++++---- openfisca_core/indexed_enums/enum_array.py | 16 +++++++++++++++- openfisca_core/indexed_enums/types.py | 9 ++------- 3 files changed, 21 insertions(+), 12 deletions(-) diff --git a/openfisca_core/indexed_enums/enum.py b/openfisca_core/indexed_enums/enum.py index 1b09a93c2..4869e5ba5 100644 --- a/openfisca_core/indexed_enums/enum.py +++ b/openfisca_core/indexed_enums/enum.py @@ -1,14 +1,12 @@ from __future__ import annotations -import enum - import numpy from . import types as t from .enum_array import EnumArray -class Enum(enum.Enum): +class Enum(t.Enum): """ Enum based on `enum34 `_, whose items have an index. @@ -19,7 +17,7 @@ def __init__(self, name: str) -> None: # When the enum item is initialized, self._member_names_ contains the # names of the previously initialized items, so its length is the index # of this item. - self.index = len(self._member_names_) + self.index = len(self._member_names_) # type: ignore[attr-defined] # Bypass the slow Enum.__eq__ __eq__ = object.__eq__ @@ -64,6 +62,7 @@ def encode( >>> encoded_array[0] 2 # Encoded value """ + if isinstance(array, EnumArray): return array @@ -95,4 +94,5 @@ def encode( [item.index for item in cls], ).astype(t.ArrayEnum) + array = numpy.asarray(array, dtype=t.ArrayEnum) return EnumArray(array, cls) diff --git a/openfisca_core/indexed_enums/enum_array.py b/openfisca_core/indexed_enums/enum_array.py index 7563ab251..de6f83df4 100644 --- a/openfisca_core/indexed_enums/enum_array.py +++ b/openfisca_core/indexed_enums/enum_array.py @@ -111,9 +111,13 @@ def decode(self) -> t.Array[t.ArrayEnum]: Decoded value: enum item """ + + if self.possible_values is None: + raise NotImplementedError + return numpy.select( [self == item.index for item in self.possible_values], - list(self.possible_values), + [item for item in self.possible_values], # type: ignore[misc] ) def decode_to_str(self) -> t.Array[t.ArrayStr]: @@ -128,6 +132,10 @@ def decode_to_str(self) -> t.Array[t.ArrayStr]: >>> enum_array.decode_to_str()[0] 'free_lodger' # String identifier """ + + if self.possible_values is None: + raise NotImplementedError + return numpy.select( [self == item.index for item in self.possible_values], [item.name for item in self.possible_values], @@ -140,12 +148,18 @@ def __str__(self) -> str: return str(self.decode_to_str()) def _is_an_enum(self, other: object) -> TypeGuard[t.Enum]: + if self.possible_values is None: + raise NotImplementedError + return ( not hasattr(other, "__name__") and other.__class__.__name__ is self.possible_values.__name__ ) def _is_an_enum_type(self, other: object) -> TypeGuard[type[t.Enum]]: + if self.possible_values is None: + raise NotImplementedError + return ( hasattr(other, "__name__") and other.__name__ is self.possible_values.__name__ diff --git a/openfisca_core/indexed_enums/types.py b/openfisca_core/indexed_enums/types.py index f8670af0f..57657b8c5 100644 --- a/openfisca_core/indexed_enums/types.py +++ b/openfisca_core/indexed_enums/types.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Protocol - from openfisca_core.types import Array as Array from openfisca_core.types import ArrayAny as ArrayAny # noqa: F401 from openfisca_core.types import ArrayBool as ArrayBool # noqa: F401 @@ -11,17 +9,14 @@ from openfisca_core.types import ArrayStr as ArrayStr # noqa: F401 import abc +import enum # Indexed enums -class Enum(Protocol): +class Enum(enum.Enum): index: int - @classmethod - def encode(cls, array: object) -> object: - ... - class EnumArray(Array[ArrayEnum], metaclass=abc.ABCMeta): ... From ebdb23b1f81c7c87606d72ad393ab543e549f4e5 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 11:43:34 +0200 Subject: [PATCH 15/20] refactor(commons): reuse array type defs --- openfisca_core/commons/formulas.py | 16 ++++++++-------- openfisca_core/commons/misc.py | 4 +--- openfisca_core/commons/rates.py | 24 ++++++++++++------------ openfisca_core/commons/types.py | 6 ++++++ 4 files changed, 27 insertions(+), 23 deletions(-) create mode 100644 openfisca_core/commons/types.py diff --git a/openfisca_core/commons/formulas.py b/openfisca_core/commons/formulas.py index bce920693..96fd0101e 100644 --- a/openfisca_core/commons/formulas.py +++ b/openfisca_core/commons/formulas.py @@ -3,14 +3,14 @@ import numpy -from openfisca_core import types as t +from . import types as t def apply_thresholds( - input: t.Array[numpy.float_], + input: t.Array[t.ArrayFloat], thresholds: t.ArrayLike[float], choices: t.ArrayLike[float], -) -> t.Array[numpy.float_]: +) -> t.Array[t.ArrayBool]: """Makes a choice based on an input and thresholds. From a list of ``choices``, this function selects one of these values @@ -39,7 +39,7 @@ def apply_thresholds( """ - condlist: list[Union[t.Array[numpy.bool_], bool]] + condlist: list[Union[t.Array[t.ArrayBool], bool]] condlist = [input <= threshold for threshold in thresholds] if len(condlist) == len(choices) - 1: @@ -58,8 +58,8 @@ def apply_thresholds( def concat( - this: Union[t.Array[numpy.str_], t.ArrayLike[str]], - that: Union[t.Array[numpy.str_], t.ArrayLike[str]], + this: Union[t.Array[t.ArrayStr], t.ArrayLike[str]], + that: Union[t.Array[t.ArrayStr], t.ArrayLike[str]], ) -> t.Array[numpy.str_]: """Concatenates the values of two arrays. @@ -89,9 +89,9 @@ def concat( def switch( - conditions: t.Array[numpy.float_], + conditions: t.Array[t.ArrayFloat], value_by_condition: Mapping[float, float], -) -> t.Array[numpy.float_]: +) -> t.Array[t.ArrayFloat]: """Mimicks a switch statement. Given an array of conditions, returns an array of the same size, diff --git a/openfisca_core/commons/misc.py b/openfisca_core/commons/misc.py index 342bbbe5f..ae66bec33 100644 --- a/openfisca_core/commons/misc.py +++ b/openfisca_core/commons/misc.py @@ -1,7 +1,5 @@ from typing import Optional, TypeVar -import numpy - from openfisca_core import types as t T = TypeVar("T") @@ -45,7 +43,7 @@ def empty_clone(original: T) -> T: return new -def stringify_array(array: Optional[t.Array[numpy.generic]]) -> str: +def stringify_array(array: Optional[t.Array[t.ArrayAny]]) -> str: """Generates a clean string representation of a numpy array. Args: diff --git a/openfisca_core/commons/rates.py b/openfisca_core/commons/rates.py index 6df1f1fee..ca34c0647 100644 --- a/openfisca_core/commons/rates.py +++ b/openfisca_core/commons/rates.py @@ -1,15 +1,15 @@ from typing import Optional -from openfisca_core.types import Array, ArrayLike - import numpy +from . import types as t + def average_rate( - target: Array[numpy.float_], - varying: ArrayLike[float], - trim: Optional[ArrayLike[float]] = None, -) -> Array[numpy.float_]: + target: t.Array[t.ArrayFloat], + varying: t.ArrayLike[float], + trim: Optional[t.ArrayLike[float]] = None, +) -> t.Array[t.ArrayFloat]: """Computes the average rate of a target net income. Given a ``target`` net income, and according to the ``varying`` gross @@ -41,7 +41,7 @@ def average_rate( """ - average_rate: Array[numpy.float_] + average_rate: t.Array[t.ArrayFloat] average_rate = 1 - target / varying @@ -62,10 +62,10 @@ def average_rate( def marginal_rate( - target: Array[numpy.float_], - varying: Array[numpy.float_], - trim: Optional[ArrayLike[float]] = None, -) -> Array[numpy.float_]: + target: t.Array[t.ArrayFloat], + varying: t.Array[t.ArrayFloat], + trim: Optional[t.ArrayLike[float]] = None, +) -> t.Array[t.ArrayFloat]: """Computes the marginal rate of a target net income. Given a ``target`` net income, and according to the ``varying`` gross @@ -97,7 +97,7 @@ def marginal_rate( """ - marginal_rate: Array[numpy.float_] + marginal_rate: t.Array[t.ArrayFloat] marginal_rate = +1 - (target[:-1] - target[1:]) / (varying[:-1] - varying[1:]) diff --git a/openfisca_core/commons/types.py b/openfisca_core/commons/types.py new file mode 100644 index 000000000..8d02a39f4 --- /dev/null +++ b/openfisca_core/commons/types.py @@ -0,0 +1,6 @@ +from openfisca_core.types import Array as Array # noqa: F401 +from openfisca_core.types import ArrayAny as ArrayAny # noqa: F401 +from openfisca_core.types import ArrayBool as ArrayBool # noqa: F401 +from openfisca_core.types import ArrayFloat as ArrayFloat # noqa: F401 +from openfisca_core.types import ArrayLike as ArrayLike # noqa: F401 +from openfisca_core.types import ArrayStr as ArrayStr # noqa: F401 From b7a74e6227fe3d67f1895a03998dd770da8c70b5 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 12:31:30 +0200 Subject: [PATCH 16/20] chore(lint): add py.typed --- openfisca_core/indexed_enums/py.typed | 0 openfisca_tasks/lint.mk | 1 + 2 files changed, 1 insertion(+) create mode 100644 openfisca_core/indexed_enums/py.typed diff --git a/openfisca_core/indexed_enums/py.typed b/openfisca_core/indexed_enums/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/openfisca_tasks/lint.mk b/openfisca_tasks/lint.mk index 7f90aa5d5..b2e83afa3 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/indexed_enums \ openfisca_core/periods \ openfisca_core/types.py @$(call print_pass,$@:) From 4a99cb4205644a04996245a995875581128fdf60 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Thu, 19 Sep 2024 12:41:00 +0200 Subject: [PATCH 17/20] chore(version): bump --- CHANGELOG.md | 6 ++++++ setup.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d3f5ff6b..a8e027023 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +### 41.5.7 [#1224](https://github.com/openfisca/openfisca-core/pull/1224) + +#### Technical changes + +- Fix doc & type definitions in the enums module + ### 41.5.6 [#1223](https://github.com/openfisca/openfisca-core/pull/1223) #### Technical changes diff --git a/setup.py b/setup.py index d62443e9e..aa7f95f43 100644 --- a/setup.py +++ b/setup.py @@ -67,7 +67,7 @@ setup( name="OpenFisca-Core", - version="41.5.6", + version="41.5.7", author="OpenFisca Team", author_email="contact@openfisca.org", classifiers=[ From f38526b6b1f6f029a9416434223d8257cb77bb86 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Mon, 16 Sep 2024 00:01:03 +0200 Subject: [PATCH 18/20] fix(types): projectors --- openfisca_core/projectors/__init__.py | 4 +- openfisca_core/projectors/helpers.py | 15 ++++--- openfisca_core/projectors/types.py | 52 ++++++++++++++++++++++++ openfisca_core/projectors/typing.py | 32 --------------- openfisca_core/simulations/simulation.py | 14 ++++--- openfisca_core/types.py | 12 +++++- openfisca_tasks/lint.mk | 1 + 7 files changed, 80 insertions(+), 50 deletions(-) create mode 100644 openfisca_core/projectors/types.py delete mode 100644 openfisca_core/projectors/typing.py 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 772a0d29e..fcb7f4898 100644 --- a/openfisca_core/types.py +++ b/openfisca_core/types.py @@ -144,13 +144,21 @@ def unit(self) -> DateUnit: # Populations -class Population(Protocol): +class CorePopulation(Protocol): + ... + + +class SinglePopulation(CorePopulation, Protocol): entity: Any def get_holder(self, variable_name: Any) -> Any: ... +class GroupPopulation(CorePopulation, Protocol): + ... + + # Simulations @@ -192,7 +200,7 @@ class Variable(Protocol): class Formula(Protocol): 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 b2e83afa3..4d1f3e097 100644 --- a/openfisca_tasks/lint.mk +++ b/openfisca_tasks/lint.mk @@ -44,6 +44,7 @@ check-types: openfisca_core/entities \ openfisca_core/indexed_enums \ openfisca_core/periods \ + openfisca_core/projectors \ openfisca_core/types.py @$(call print_pass,$@:) From 4e3dfcb00b1830a2a85e851844a3e8050fa955d0 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Tue, 17 Sep 2024 01:28:25 +0200 Subject: [PATCH 19/20] chore(version): bump --- CHANGELOG.md | 6 ++++++ setup.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a8e027023..609589693 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +### 41.5.8 [#1220](https://github.com/openfisca/openfisca-core/pull/1221) + +#### Technical changes + +- Fix doc & type definitions in the projectors module + ### 41.5.7 [#1224](https://github.com/openfisca/openfisca-core/pull/1224) #### Technical changes diff --git a/setup.py b/setup.py index aa7f95f43..9aa018fca 100644 --- a/setup.py +++ b/setup.py @@ -67,7 +67,7 @@ setup( name="OpenFisca-Core", - version="41.5.7", + version="41.5.8", author="OpenFisca Team", author_email="contact@openfisca.org", classifiers=[ From 584e97fd4caea9ecb173494ad81a62adebd3a8e5 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Wed, 18 Sep 2024 13:43:24 +0200 Subject: [PATCH 20/20] doc(types): add projector types --- .../projectors/entity_to_person_projector.py | 2 +- .../first_person_to_entity_projector.py | 2 +- openfisca_core/projectors/helpers.py | 18 ++++++------- openfisca_core/projectors/types.py | 25 +++++++------------ .../unique_role_to_entity_projector.py | 2 +- 5 files changed, 20 insertions(+), 29 deletions(-) diff --git a/openfisca_core/projectors/entity_to_person_projector.py b/openfisca_core/projectors/entity_to_person_projector.py index ca6245a1f..392fda08a 100644 --- a/openfisca_core/projectors/entity_to_person_projector.py +++ b/openfisca_core/projectors/entity_to_person_projector.py @@ -4,7 +4,7 @@ class EntityToPersonProjector(Projector): """For instance person.family.""" - def __init__(self, entity, parent=None): + def __init__(self, entity, parent=None) -> None: self.reference_entity = entity self.parent = parent diff --git a/openfisca_core/projectors/first_person_to_entity_projector.py b/openfisca_core/projectors/first_person_to_entity_projector.py index 4b4e7b799..d986460cd 100644 --- a/openfisca_core/projectors/first_person_to_entity_projector.py +++ b/openfisca_core/projectors/first_person_to_entity_projector.py @@ -4,7 +4,7 @@ class FirstPersonToEntityProjector(Projector): """For instance famille.first_person.""" - def __init__(self, entity, parent=None): + def __init__(self, entity, parent=None) -> None: self.target_entity = entity self.reference_entity = entity.members self.parent = parent diff --git a/openfisca_core/projectors/helpers.py b/openfisca_core/projectors/helpers.py index 9c666eb6b..d3d8a21c7 100644 --- a/openfisca_core/projectors/helpers.py +++ b/openfisca_core/projectors/helpers.py @@ -4,7 +4,7 @@ from openfisca_core import entities, projectors -from .types import GroupEntity, GroupPopulation, Role, SingleEntity, SinglePopulation +from . import types as t def projectable(function): @@ -17,10 +17,10 @@ def projectable(function): def get_projector_from_shortcut( - population: SinglePopulation | GroupPopulation, + population: t.CorePopulation, shortcut: str, - parent: projectors.Projector | None = None, -) -> projectors.Projector | None: + parent: t.Projector | None = None, +) -> t.Projector | None: """Get a projector from a shortcut. Projectors are used to project an invidividual Population's or a @@ -108,12 +108,10 @@ def get_projector_from_shortcut( """ - entity: SingleEntity | GroupEntity = population.entity + entity: t.CoreEntity = population.entity if isinstance(entity, entities.Entity): - populations: Mapping[ - str, SinglePopulation | GroupPopulation - ] = population.simulation.populations + populations: Mapping[str, t.CorePopulation] = population.simulation.populations if shortcut not in populations.keys(): return None @@ -124,8 +122,8 @@ def get_projector_from_shortcut( return projectors.FirstPersonToEntityProjector(population, parent) if isinstance(entity, entities.GroupEntity): - roles: Iterable[Role] = entity.roles - role: Role | None = entities.find_role(roles, shortcut, total=1) + roles: Iterable[t.Role] = entity.roles + role: t.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 index a090c4813..e145cc9a1 100644 --- a/openfisca_core/projectors/types.py +++ b/openfisca_core/projectors/types.py @@ -8,11 +8,7 @@ # Entities -class SingleEntity(t.SingleEntity, Protocol): - ... - - -class GroupEntity(t.GroupEntity, Protocol): +class CoreEntity(t.CoreEntity, Protocol): ... @@ -20,22 +16,19 @@ class Role(t.Role, Protocol): ... -# Populations +# Projectors -class SinglePopulation(t.SinglePopulation, Protocol): - @property - def entity(self) -> t.SingleEntity: - ... +class Projector(Protocol): + ... - @property - def simulation(self) -> Simulation: - ... + +# Populations -class GroupPopulation(t.GroupPopulation, Protocol): +class CorePopulation(t.CorePopulation, Protocol): @property - def entity(self) -> t.GroupEntity: + def entity(self) -> CoreEntity: ... @property @@ -48,5 +41,5 @@ def simulation(self) -> Simulation: class Simulation(t.Simulation, Protocol): @property - def populations(self) -> Mapping[str, SinglePopulation | GroupPopulation]: + def populations(self) -> Mapping[str, CorePopulation]: ... diff --git a/openfisca_core/projectors/unique_role_to_entity_projector.py b/openfisca_core/projectors/unique_role_to_entity_projector.py index fed2f249c..c56548433 100644 --- a/openfisca_core/projectors/unique_role_to_entity_projector.py +++ b/openfisca_core/projectors/unique_role_to_entity_projector.py @@ -4,7 +4,7 @@ class UniqueRoleToEntityProjector(Projector): """For instance famille.declarant_principal.""" - def __init__(self, entity, role, parent=None): + def __init__(self, entity, role, parent=None) -> None: self.target_entity = entity self.reference_entity = entity.members self.parent = parent