X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=datetime_utils.py;fp=datetime_utils.py;h=60b859afd0a0ccfec361034f66b2bff615e5bc43;hb=e70297b7e29a07457d6c7195425e734a1290abeb;hp=3565936fce66c1197a04a8926f902452e6350ac4;hpb=679b41526d9327f5bcecc702a928a72b68a1e0c8;p=python_utils.git

diff --git a/datetime_utils.py b/datetime_utils.py
index 3565936..60b859a 100644
--- a/datetime_utils.py
+++ b/datetime_utils.py
@@ -304,6 +304,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:
@@ -364,16 +375,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: