Various

[python_utils.git] / datetime_utils.py

diff --git a/datetime_utils.py b/datetime_utils.py
index 0b94283b01df595300ecc448f108303c2dba2b2e..7787c6f0b0b74b84b84fca28bba8d2381e1d8668 100644 (file)
--- 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
@@ -51,12 +51,18 @@ def date_and_time_to_datetime(date: datetime.date,
     )
 
 
-def datetime_to_date(date: datetime.datetime) -> datetime.date:
-    return datetime.date(
-        date.year,
-        date.month,
-        date.day
-    )
+def datetime_to_date_and_time(
+        dt: datetime.datetime
+) -> Tuple[datetime.date, datetime.time]:
+    return (dt.date(), dt.timetz())
+
+
+def datetime_to_date(dt: datetime.datetime) -> datetime.date:
+    return datetime_to_date_and_time(dt)[0]
+
+
+def datetime_to_time(dt: datetime.datetime) -> datetime.time:
+    return datetime_to_date_and_time(dt)[1]
 
 
 # An enum to represent units with which we can compute deltas.
@@ -77,12 +83,25 @@ class TimeUnit(enum.Enum):
     MONTHS = 13
     YEARS = 14
 
+    @classmethod
+    def is_valid(cls, value: Any):
+        if type(value) is int:
+            return value in cls._value2member_map_
+        elif type(value) is TimeUnit:
+            return value.value in cls._value2member_map_
+        elif type(value) is 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
 
@@ -317,6 +336,8 @@ def minute_number_to_time_string(minute_num: MinuteOfDay) -> str:
 
 def parse_duration(duration: str) -> int:
     """Parse a duration in string form."""
+    if duration.isdigit():
+        return int(duration)
     seconds = 0
     m = re.search(r'(\d+) *d[ays]*', duration)
     if m is not None: