Add add_timezone.

diff --git a/datetime_utils.py b/datetime_utils.py
index fc1ca4ed9e07e407554da7e300f53ca58f98661a..cbebf4dffada4955703859186d3b4d2d411dec53 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 Any, NewType, Tuple
+from typing import Any, NewType, Optional, Tuple
 
 import holidays  # type: ignore
 import pytz
@@ -31,13 +31,92 @@ def is_timezone_aware(dt: datetime.datetime) -> bool:
 
 
 def is_timezone_naive(dt: datetime.datetime) -> bool:
+    """Inverse of is_timezone_aware.
+
+    >>> is_timezone_naive(datetime.datetime.now())
+    True
+
+    >>> is_timezone_naive(now_pacific())
+    False
+
+    """
     return not is_timezone_aware(dt)
 
 
-def replace_timezone(dt: datetime.datetime, tz: datetime.tzinfo) -> datetime.datetime:
+def strip_timezone(dt: datetime.datetime) -> datetime.datetime:
+    """Remove the timezone from a datetime.  Does not change the
+    hours, minutes, seconds, months, days, years, etc... Thus the
+    instant to which this timestamp refers will change.  Silently
+    ignores datetimes which are already timezone naive.
+
+    >>> now = now_pacific()
+    >>> now.tzinfo == None
+    False
+
+    >>> dt = strip_timezone(now)
+    >>> dt == now
+    False
+
+    >>> dt.tzinfo == None
+    True
+
+    >>> dt.hour == now.hour
+    True
+
     """
-    Replaces the timezone on a datetime object directly (leaving
-    the year, month, day, hour, minute, second, micro, etc... alone).
+    if is_timezone_naive(dt):
+        return dt
+    return replace_timezone(dt, None)
+
+
+def add_timezone(dt: datetime.datetime, tz: datetime.tzinfo) -> datetime.datetime:
+    """
+    Adds a timezone to a timezone naive datetime.  This does not
+    change the instant to which the timestamp refers.  See also:
+    replace_timezone.
+
+    >>> now = datetime.datetime.now()
+    >>> is_timezone_aware(now)
+    False
+
+    >>> now_pacific = add_timezone(now, pytz.timezone('US/Pacific'))
+    >>> is_timezone_aware(now_pacific)
+    True
+
+    >>> now.hour == now_pacific.hour
+    True
+    >>> now.minute == now_pacific.minute
+    True
+
+    """
+
+    # This doesn't work, tz requires a timezone naive dt.  Two options
+    # here:
+    #     1. Use strip_timezone and try again.
+    #     2. Replace the timezone on your dt object via replace_timezone.
+    #        Be aware that this changes the instant to which the dt refers
+    #        and, further, can introduce weirdness like UTC offsets that
+    #        are weird (e.g. not an even multiple of an hour, etc...)
+    if is_timezone_aware(dt):
+        if dt.tzinfo == tz:
+            return dt
+        raise Exception(
+            f'{dt} is already timezone aware; use replace_timezone or translate_timezone '
+            + 'depending on the semantics you want.'
+        )
+    return dt.replace(tzinfo=tz)
+
+
+def replace_timezone(dt: datetime.datetime, tz: Optional[datetime.tzinfo]) -> datetime.datetime:
+    """Replaces the timezone on a timezone aware datetime object directly
+    (leaving the year, month, day, hour, minute, second, micro,
+    etc... alone).
+
+    Works with timezone aware and timezone naive dts but for the
+    latter it is probably better to use add_timezone or just create it
+    with a tz parameter.  Using this can have weird side effects like
+    UTC offsets that are not an even multiple of an hour, etc...
+
     Note: this changes the instant to which this dt refers.
 
     >>> from pytz import UTC
@@ -52,22 +131,33 @@ def replace_timezone(dt: datetime.datetime, tz: datetime.tzinfo) -> datetime.dat
     True
 
     """
-    return datetime.datetime(
-        dt.year,
-        dt.month,
-        dt.day,
-        dt.hour,
-        dt.minute,
-        dt.second,
-        dt.microsecond,
-        tzinfo=tz,
-    )
+    if is_timezone_aware(dt):
+        logger.warning(
+            '%s already has a timezone; klobbering it anyway.\n  Be aware that this operation changed the instant to which the object refers.',
+            dt,
+        )
+        return datetime.datetime(
+            dt.year,
+            dt.month,
+            dt.day,
+            dt.hour,
+            dt.minute,
+            dt.second,
+            dt.microsecond,
+            tzinfo=tz,
+        )
+    else:
+        if tz:
+            return add_timezone(dt, tz)
+        else:
+            return dt
 
 
 def replace_time_timezone(t: datetime.time, tz: datetime.tzinfo) -> datetime.time:
     """
     Replaces the timezone on a datetime.time directly without performing
-    any translation.
+    any translation.  Note that, as above, this will change the instant
+    to which the time refers.
 
     >>> t = datetime.time(8, 15, 12, 0, pytz.UTC)
     >>> t.tzname()
@@ -104,7 +194,7 @@ def translate_timezone(dt: datetime.datetime, tz: datetime.tzinfo) -> datetime.d
 
 def now() -> datetime.datetime:
     """
-    What time is it?  Result returned in UTC
+    What time is it?  Returned as a timezone naive datetime.
     """
     return datetime.datetime.now()