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