ACL uses enums, some more tests, other stuff.

1 files changed, 17 insertions, 1 deletions

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