1818from typing import Dict , Optional , Union
1919
2020from undate .converters .base import BaseDateConverter
21- from undate .date import ONE_DAY , ONE_MONTH_MAX , Date , DatePrecision , Timedelta
21+ from undate .date import ONE_DAY , ONE_MONTH_MAX , Date , DatePrecision , Timedelta , Udelta
2222
2323
2424class Calendar (StrEnum ):
@@ -420,13 +420,14 @@ def _get_date_part(self, part: str) -> Optional[str]:
420420 value = self .initial_values .get (part )
421421 return str (value ) if value else None
422422
423- def duration (self ) -> Timedelta :
423+ def duration (self ) -> Timedelta | Udelta :
424424 """What is the duration of this date?
425425 Calculate based on earliest and latest date within range,
426426 taking into account the precision of the date even if not all
427427 parts of the date are known. Note that durations are inclusive
428428 (i.e., a closed interval) and include both the earliest and latest
429- date rather than the difference between them."""
429+ date rather than the difference between them. Returns a :class:`undate.date.Timedelta` when
430+ possible, and an :class:`undate.date.Udelta` when the duration is uncertain."""
430431
431432 # if precision is a single day, duration is one day
432433 # no matter when it is or what else is known
@@ -437,24 +438,48 @@ def duration(self) -> Timedelta:
437438 # calculate month duration within a single year (not min/max)
438439 if self .precision == DatePrecision .MONTH :
439440 latest = self .latest
441+ # if year is unknown, calculate month duration in
442+ # leap year and non-leap year, in case length varies
440443 if not self .known_year :
441- # if year is unknown, calculate month duration in
442- # a single year
443- latest = Date (self .earliest .year , self .latest .month , self .latest .day )
444+ # TODO: should leap-year specific logic shift to the calendars,
445+ # since it works differently depending on the calendar?
446+ possible_years = [
447+ self .calendar_converter .LEAP_YEAR ,
448+ self .calendar_converter .NON_LEAP_YEAR ,
449+ ]
450+ # TODO: what about partially known years like 191X ?
451+ else :
452+ # otherwise, get possible durations for all possible months
453+ # for a known year
454+ possible_years = [self .earliest .year ]
455+
456+ # for every possible month and year, get max days for that month,
457+ possible_max_days = set ()
458+ # appease mypy, which says month values could be None here
459+ if self .earliest .month is not None and self .latest .month is not None :
460+ for possible_month in range (self .earliest .month , self .latest .month + 1 ):
461+ for year in possible_years :
462+ possible_max_days .add (
463+ self .calendar_converter .max_day (year , possible_month )
464+ )
465+
466+ # if there is more than one possible value for month length,
467+ # whether due to leap year / non-leap year or ambiguous month,
468+ # return a uncertain delta
469+ if len (possible_max_days ) > 1 :
470+ return Udelta (* possible_max_days )
471+
472+ # otherwise, calculate timedelta normally
473+ max_day = list (possible_max_days )[0 ]
474+ latest = Date (self .earliest .year , self .earliest .month , max_day )
444475
445- # TODO: calculate duration for a leap year and a non-leap year,
446- # then return a udelta if they vary
447- # TODO: how does this logic work for other calendars?
448-
449- # latest = datetime.date(
450- # self.earliest.year, self.latest.month, self.latest.day
451- # )
452476 delta = latest - self .earliest + ONE_DAY
453477 # month duration can't ever be more than 31 days
454478 # (could we ever know if it's smaller?)
455479
456480 # if granularity == month but not known month, duration = 31
457481 if delta .astype (int ) > 31 :
482+ # FIXME: this depends on calendar!
458483 return ONE_MONTH_MAX
459484 return delta
460485
0 commit comments