my.time.tz: implement different policies for localizing

my/core/common.py

@@ -273,7 +273,8 @@ class classproperty(Generic[_R]):
 #     def __get__(self) -> _R:
 #         return self.f()
 
-# TODO maybe use opaque mypy alias?
+# for now just serves documentation purposes... but one day might make it statically verifiable where possible?
+# TODO e.g. maybe use opaque mypy alias?
 tzdatetime = datetime
 
 

my/core/error.py

@@ -15,6 +15,11 @@ ResT = Union[T, E]
 Res = ResT[T, Exception]
 
 
+def notnone(x: Optional[T]) -> T:
+    assert x is not None
+    return x
+
+
 def unwrap(res: Res[T]) -> T:
     if isinstance(res, Exception):
         raise res

my/time/tz/common.py

@@ -0,0 +1,41 @@
+from datetime import datetime
+from typing import Callable, cast
+
+from ...core.common import tzdatetime, Literal
+
+
+'''
+Depending on the specific data provider and your level of paranoia you might expect different behaviour.. E.g.:
+- if your objects already have tz info, you might not need to call localize() at all
+- it's safer when either all of your objects are tz aware or all are tz unware, not a mixture
+- you might trust your original timezone, or it might just be UTC, and you want to use something more reasonable
+'''
+Policy = Literal[
+    'keep'   , # if datetime is tz aware, just preserve it
+    'convert', # if datetime is tz aware, convert to provider's tz
+    'throw'  , # if datetime is tz aware, throw exception
+    # todo 'warn'? not sure if very useful
+]
+
+def default_policy() -> Policy:
+    try:
+        from my.config import time as user_config
+        return cast(Policy, user_config.tz.policy)
+    except Exception as e:
+        # todo meh.. need to think how to do this more carefully
+        # rationale: do not mess with user's data until they want
+        return 'keep'
+
+
+def localize_with_policy(lfun: Callable[[datetime], tzdatetime], dt: datetime, policy: Policy=default_policy()) -> tzdatetime:
+    tz = dt.tzinfo
+    if tz is None:
+        return lfun(dt)
+
+    if   policy == 'keep':
+        return dt
+    elif policy == 'convert':
+        ldt = lfun(dt.replace(tzinfo=None))
+        return dt.astimezone(ldt.tzinfo)
+    else: # policy == 'error':
+        raise RuntimeError(f"{dt} already has timezone information (use 'policy' argument to adjust this behaviour)")

my/time/tz/main.py

@@ -2,8 +2,12 @@
 Timezone data provider
 '''
 from datetime import datetime
+from ...core.common import tzdatetime
 
-def localize(dt: datetime) -> datetime:
-    # For now, it's user's reponsibility to check that it actually managed to localize
+# todo hmm, kwargs isn't mypy friendly.. but specifying types would require duplicating default args. uhoh
+def localize(dt: datetime, **kwargs) -> tzdatetime:
+    # todo document patterns for combining multiple data sources
+    # e.g. see https://github.com/karlicoss/HPI/issues/89#issuecomment-716495136
     from . import via_location as L
-    return L.localize(dt)
+    from .common import localize_with_policy
+    return localize_with_policy(L.localize, dt, **kwargs)

my/time/tz/via_location.py

@@ -17,7 +17,7 @@ from typing import Dict, Iterator, List, NamedTuple, Optional, Tuple
 from more_itertools import seekable
 import pytz
 
-from ...core.common import LazyLogger, mcachew
+from ...core.common import LazyLogger, mcachew, tzdatetime
 from ...core.cachew import cache_dir
 from ...location.google import locations
 
@@ -130,11 +130,10 @@ def _get_tz(dt: datetime) -> Optional[pytz.BaseTzInfo]:
     return _get_home_tz(loc=loc)
 
 
-def localize(dt: datetime) -> datetime:
-    # todo not sure. warn instead?
-    assert dt.tzinfo is None, dt
+def localize(dt: datetime) -> tzdatetime:
     tz = _get_tz(dt)
     if tz is None:
+        # TODO -- this shouldn't really happen.. think about it carefully later
         return dt
     else:
         return tz.localize(dt)