ACL uses enums, some more tests, other stuff.
[python_utils.git] / datetime_utils.py
index d70bf4a79008effe5e2a3aec8684d3352be6c78d..795b427c31b4e57a4551aab26badd1f81b68c9a2 100644 (file)
@@ -3,10 +3,12 @@
 """Utilities related to dates and times and datetimes."""
 
 import datetime
+import enum
 import logging
 import re
-from typing import NewType
+from typing import Any, NewType, Tuple
 
+import holidays  # type: ignore
 import pytz
 
 import constants
@@ -14,29 +16,189 @@ import constants
 logger = logging.getLogger(__name__)
 
 
-def now_pst() -> datetime.datetime:
-    return datetime.datetime.now(tz=pytz.timezone("US/Pacific"))
+def replace_timezone(dt: datetime.datetime,
+                     tz: datetime.tzinfo) -> datetime.datetime:
+    return dt.replace(tzinfo=None).astimezone(tz=tz)
 
 
 def now() -> datetime.datetime:
     return datetime.datetime.now()
 
 
-def datetime_to_string(
-    dt: datetime.datetime,
-    *,
-    date_time_separator=" ",
-    include_timezone=True,
-    include_dayname=False,
-    include_seconds=True,
-    include_fractional=False,
-    twelve_hour=True,
+def now_pst() -> datetime.datetime:
+    return replace_timezone(now(), pytz.timezone("US/Pacific"))
+
+
+def date_to_datetime(date: datetime.date) -> datetime.datetime:
+    return datetime.datetime(
+        date.year,
+        date.month,
+        date.day,
+        0, 0, 0, 0
+    )
+
+
+def date_and_time_to_datetime(date: datetime.date,
+                              time: datetime.time) -> datetime.datetime:
+    return datetime.datetime(
+        date.year,
+        date.month,
+        date.day,
+        time.hour,
+        time.minute,
+        time.second,
+        time.millisecond
+    )
+
+
+def datetime_to_date(date: datetime.datetime) -> datetime.date:
+    return datetime.date(
+        date.year,
+        date.month,
+        date.day
+    )
+
+
+# An enum to represent units with which we can compute deltas.
+class TimeUnit(enum.Enum):
+    MONDAYS = 0
+    TUESDAYS = 1
+    WEDNESDAYS = 2
+    THURSDAYS = 3
+    FRIDAYS = 4
+    SATURDAYS = 5
+    SUNDAYS = 6
+    SECONDS = 7
+    MINUTES = 8
+    HOURS = 9
+    DAYS = 10
+    WORKDAYS = 11
+    WEEKS = 12
+    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
+
+    # N days from base
+    if unit == TimeUnit.DAYS:
+        timedelta = datetime.timedelta(days=count)
+        return base + timedelta
+
+    # N workdays from base
+    elif unit == TimeUnit.WORKDAYS:
+        if count < 0:
+            count = abs(count)
+            timedelta = datetime.timedelta(days=-1)
+        else:
+            timedelta = datetime.timedelta(days=1)
+        skips = holidays.US(years=base.year).keys()
+        while count > 0:
+            old_year = base.year
+            base += timedelta
+            if base.year != old_year:
+                skips = holidays.US(years=base.year).keys()
+            if (
+                    base.weekday() < 5 and
+                    datetime.date(base.year,
+                                  base.month,
+                                  base.day) not in skips
+            ):
+                count -= 1
+        return base
+
+    # N weeks from base
+    elif unit == TimeUnit.WEEKS:
+        timedelta = datetime.timedelta(weeks=count)
+        base = base + timedelta
+        return base
+
+    # N months from base
+    elif unit == TimeUnit.MONTHS:
+        month_term = count % 12
+        year_term = count // 12
+        new_month = base.month + month_term
+        if new_month > 12:
+            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,
+        )
+
+    # N years from base
+    elif unit == TimeUnit.YEARS:
+        new_year = base.year + count
+        return datetime.datetime(
+            new_year,
+            base.month,
+            base.day,
+            base.hour,
+            base.minute,
+            base.second,
+            base.microsecond,
+        )
+
+    # N weekdays from base (e.g. 4 wednesdays from today)
+    direction = 1 if count > 0 else -1
+    count = abs(count)
+    timedelta = datetime.timedelta(days=direction)
+    start = base
+    while True:
+        dow = base.weekday()
+        if dow == unit and start != base:
+            count -= 1
+            if count == 0:
+                return base
+        base = base + timedelta
+
+
+def get_format_string(
+        *,
+        date_time_separator=" ",
+        include_timezone=True,
+        include_dayname=False,
+        use_month_abbrevs=False,
+        include_seconds=True,
+        include_fractional=False,
+        twelve_hour=True,
 ) -> str:
-    """A nice way to convert a datetime into a string."""
     fstring = ""
     if include_dayname:
         fstring += "%a/"
-    fstring = f"%Y/%b/%d{date_time_separator}"
+
+    if use_month_abbrevs:
+        fstring = f"%Y/%b/%d{date_time_separator}"
+    else:
+        fstring = f"%Y/%m/%d{date_time_separator}"
     if twelve_hour:
         fstring += "%I:%M"
         if include_seconds:
@@ -50,9 +212,58 @@ def datetime_to_string(
         fstring += ".%f"
     if include_timezone:
         fstring += "%z"
+    return fstring
+
+
+def datetime_to_string(
+    dt: datetime.datetime,
+    *,
+    date_time_separator=" ",
+    include_timezone=True,
+    include_dayname=False,
+    use_month_abbrevs=False,
+    include_seconds=True,
+    include_fractional=False,
+    twelve_hour=True,
+) -> str:
+    """A nice way to convert a datetime into a string."""
+    fstring = get_format_string(
+        date_time_separator=date_time_separator,
+        include_timezone=include_timezone,
+        include_dayname=include_dayname,
+        include_seconds=include_seconds,
+        include_fractional=include_fractional,
+        twelve_hour=twelve_hour)
     return dt.strftime(fstring).strip()
 
 
+def string_to_datetime(
+        txt: str,
+        *,
+        date_time_separator=" ",
+        include_timezone=True,
+        include_dayname=False,
+        use_month_abbrevs=False,
+        include_seconds=True,
+        include_fractional=False,
+        twelve_hour=True,
+) -> Tuple[datetime.datetime, str]:
+    """A nice way to convert a string into a datetime.  Also consider
+    dateparse.dateparse_utils for a full parser.
+    """
+    fstring = get_format_string(
+        date_time_separator=date_time_separator,
+        include_timezone=include_timezone,
+        include_dayname=include_dayname,
+        include_seconds=include_seconds,
+        include_fractional=include_fractional,
+        twelve_hour=twelve_hour)
+    return (
+        datetime.datetime.strptime(txt, fstring),
+        fstring
+    )
+
+
 def timestamp() -> str:
     """Return a timestamp for now in Pacific timezone."""
     ts = datetime.datetime.now(tz=pytz.timezone("US/Pacific"))
@@ -104,7 +315,9 @@ def datetime_to_minute_number(dt: datetime.datetime) -> MinuteOfDay:
 
 
 def minute_number_to_time_string(minute_num: MinuteOfDay) -> str:
-    """Convert minute number from start of day into hour:minute am/pm string."""
+    """Convert minute number from start of day into hour:minute am/pm
+    string.
+    """
     hour = minute_num // 60
     minute = minute_num % 60
     ampm = "a"
@@ -171,5 +384,6 @@ def describe_duration_briefly(age: int) -> str:
         descr = f"{int(days[0])}d "
     if hours[0] > 0:
         descr = descr + f"{int(hours[0])}h "
-    descr = descr + f"{int(minutes[0])}m"
-    return descr
+    if minutes[0] > 0 or len(descr) == 0:
+        descr = descr + f"{int(minutes[0])}m"
+    return descr.strip()