From 09e6d10face80d98a4578ff54192b5c8bec007d7 Mon Sep 17 00:00:00 2001
From: Scott Gasch <scott@gasch.org>
Date: Fri, 9 Jul 2021 18:31:02 -0700
Subject: ACL uses enums, some more tests, other stuff.

---
 datetime_utils.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

(limited to 'datetime_utils.py')

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
 
-- 
cgit v1.3