X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=datetime_utils.py;h=795b427c31b4e57a4551aab26badd1f81b68c9a2;hb=09e6d10face80d98a4578ff54192b5c8bec007d7;hp=0b94283b01df595300ecc448f108303c2dba2b2e;hpb=11eeb8574b7b4620ac6fd440cb251f8aa2458f5b;p=python_utils.git

diff --git a/datetime_utils.py b/datetime_utils.py
index 0b94283..795b427 100644
--- a/datetime_utils.py
+++ b/datetime_utils.py
@@ -6,7 +6,7 @@ import datetime
 import enum
 import logging
 import re
-from typing import NewType, Tuple
+from typing import Any, NewType, Tuple
 
 import holidays  # type: ignore
 import pytz
@@ -77,12 +77,28 @@ class TimeUnit(enum.Enum):
     MONTHS = 13
     YEARS = 14
 
+    @classmethod
+    def is_valid(cls, value: Any):
+        if type(value) is int:
+            print("int")
+            return value in cls._value2member_map_
+        elif type(value) is TimeUnit:
+            print("TimeUnit")
+            return value.value in cls._value2member_map_
+        elif type(value) is str:
+            print("str")
+            return value in cls._member_names_
+        else:
+            print(type(value))
+            return False
+
 
 def n_timeunits_from_base(
     count: int,
     unit: TimeUnit,
     base: datetime.datetime
 ) -> datetime.datetime:
+    assert TimeUnit.is_valid(unit)
     if count == 0:
         return base