X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=datetime_utils.py;h=1cee5163a22179d3c634a078433e83c53add8109;hb=b22b39493c5b6c747b16e9430f3833bb8869cef6;hp=97947203f6a126348d1b6602c320b07909f83970;hpb=36fea7f15ed17150691b5b3ead75450e575229ef;p=python_utils.git

diff --git a/datetime_utils.py b/datetime_utils.py
index 9794720..1cee516 100644
--- a/datetime_utils.py
+++ b/datetime_utils.py
@@ -34,9 +34,7 @@ def is_timezone_naive(dt: datetime.datetime) -> bool:
     return not is_timezone_aware(dt)
 
 
-def replace_timezone(
-    dt: datetime.datetime, tz: datetime.tzinfo
-) -> datetime.datetime:
+def replace_timezone(dt: datetime.datetime, tz: datetime.tzinfo) -> datetime.datetime:
     """
     Replaces the timezone on a datetime object directly (leaving
     the year, month, day, hour, minute, second, micro, etc... alone).
@@ -66,9 +64,7 @@ def replace_timezone(
     )
 
 
-def replace_time_timezone(
-    t: datetime.time, tz: datetime.tzinfo
-) -> datetime.time:
+def replace_time_timezone(t: datetime.time, tz: datetime.tzinfo) -> datetime.time:
     """
     Replaces the timezone on a datetime.time directly without performing
     any translation.
@@ -85,9 +81,7 @@ def replace_time_timezone(
     return t.replace(tzinfo=tz)
 
 
-def translate_timezone(
-    dt: datetime.datetime, tz: datetime.tzinfo
-) -> datetime.datetime:
+def translate_timezone(dt: datetime.datetime, tz: datetime.tzinfo) -> datetime.datetime:
     """
     Translates dt into a different timezone by adjusting the year, month,
     day, hour, minute, second, micro, etc... appropriately.  The returned
@@ -170,9 +164,7 @@ def time_to_datetime_today(time: datetime.time) -> datetime.datetime:
     return datetime.datetime.combine(now, time, tz)
 
 
-def date_and_time_to_datetime(
-    date: datetime.date, time: datetime.time
-) -> datetime.datetime:
+def date_and_time_to_datetime(date: datetime.date, time: datetime.time) -> datetime.datetime:
     """
     Given a date and time, merge them and return a datetime.
 
@@ -235,7 +227,7 @@ def datetime_to_time(dt: datetime.datetime) -> datetime.time:
     return datetime_to_date_and_time(dt)[1]
 
 
-class TimeUnit(enum.Enum):
+class TimeUnit(enum.IntEnum):
     """An enum to represent units with which we can compute deltas."""
 
     MONDAYS = 0
@@ -267,9 +259,7 @@ class TimeUnit(enum.Enum):
             return False
 
 
-def n_timeunits_from_base(
-    count: int, unit: TimeUnit, base: datetime.datetime
-) -> datetime.datetime:
+def n_timeunits_from_base(count: int, unit: TimeUnit, base: datetime.datetime) -> datetime.datetime:
     """Return a datetime that is N units before/after a base datetime.
     e.g.  3 Wednesdays from base datetime, 2 weeks from base date, 10
     years before base datetime, 13 minutes after base datetime, etc...
@@ -310,6 +300,17 @@ def n_timeunits_from_base(
     >>> n_timeunits_from_base(50, TimeUnit.SECONDS, base)
     datetime.datetime(2021, 9, 10, 11, 25, 41, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=61200)))
 
+    Next month corner case -- it will try to make Feb 31, 2022 then count
+    backwards.
+    >>> base = string_to_datetime("2022/01/31 11:24:51AM-0700")[0]
+    >>> n_timeunits_from_base(1, TimeUnit.MONTHS, base)
+    datetime.datetime(2022, 2, 28, 11, 24, 51, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=61200)))
+
+    Last month with the same corner case
+    >>> base = string_to_datetime("2022/03/31 11:24:51AM-0700")[0]
+    >>> n_timeunits_from_base(-1, TimeUnit.MONTHS, base)
+    datetime.datetime(2022, 2, 28, 11, 24, 51, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=61200)))
+
     """
     assert TimeUnit.is_valid(unit)
     if count == 0:
@@ -348,10 +349,7 @@ def n_timeunits_from_base(
             base += timedelta
             if base.year != old_year:
                 skips = holidays.US(years=base.year).keys()
-            if (
-                base.weekday() < 5
-                and datetime.date(base.year, base.month, base.day) not in skips
-            ):
+            if base.weekday() < 5 and datetime.date(base.year, base.month, base.day) not in skips:
                 count -= 1
         return base
 
@@ -370,16 +368,23 @@ def n_timeunits_from_base(
             new_month %= 12
             year_term += 1
         new_year = base.year + year_term
-        return datetime.datetime(
-            new_year,
-            new_month,
-            base.day,
-            base.hour,
-            base.minute,
-            base.second,
-            base.microsecond,
-            base.tzinfo,
-        )
+        day = base.day
+        while True:
+            try:
+                ret = datetime.datetime(
+                    new_year,
+                    new_month,
+                    day,
+                    base.hour,
+                    base.minute,
+                    base.second,
+                    base.microsecond,
+                    base.tzinfo,
+                )
+                break
+            except ValueError:
+                day -= 1
+        return ret
 
     # N years from base
     elif unit == TimeUnit.YEARS:
@@ -768,7 +773,7 @@ def describe_timedelta(delta: datetime.timedelta) -> str:
     '1 day, and 10 minutes'
 
     """
-    return describe_duration(delta.total_seconds())
+    return describe_duration(int(delta.total_seconds()))  # Note: drops milliseconds
 
 
 def describe_duration_briefly(seconds: int, *, include_seconds=False) -> str:
@@ -813,7 +818,7 @@ def describe_timedelta_briefly(delta: datetime.timedelta) -> str:
     '1d 10m'
 
     """
-    return describe_duration_briefly(delta.total_seconds())
+    return describe_duration_briefly(int(delta.total_seconds()))  # Note: drops milliseconds
 
 
 if __name__ == '__main__':