my.time.tz: implement different policies for localizing

doc/example_config/my/config/__init__.py

@@ -42,3 +42,9 @@ class location:
     # todo ugh, need to think about it... mypy wants the type here to be general, otherwise it can't deduce
     # and we can't import the types from the module itself, otherwise would be circular. common module?
     home: Union[LatLon, Sequence[Tuple[DateIsh, LatLon]]] = (1.0, -1.0)
+
+
+# todo hmm it's getting out of hand.. perhaps better to keep stubs in the actual my.config presetn in the repository instead
+class time:
+    class tz:
+        pass

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)

tests/tz.py

@@ -3,7 +3,9 @@ from pathlib import Path
 import sys
 
 import pytest # type: ignore
+import pytz # type: ignore
 
+from my.core.error import notnone
 import my.time.tz.main as TZ
 import my.time.tz.via_location as LTZ
 
@@ -33,24 +35,38 @@ def test_tz() -> None:
 
     # not present in the test data
     tz = LTZ._get_tz(D('20200101 10:00:00'))
-    assert tz is not None
-    assert tz.zone == 'Europe/Sofia'
+    assert notnone(tz).zone == 'Europe/Sofia'
 
     tz = LTZ._get_tz(D('20170801 11:00:00'))
-    assert tz is not None
-    assert tz.zone == 'Europe/Vienna'
+    assert notnone(tz).zone == 'Europe/Vienna'
 
     tz = LTZ._get_tz(D('20170730 10:00:00'))
-    assert tz is not None
-    assert tz.zone == 'Europe/Rome'
+    assert notnone(tz).zone == 'Europe/Rome'
 
     tz = LTZ._get_tz(D('20201001 14:15:16'))
     assert tz is not None
-    assert tz.zone == 'Europe/Moscow'
 
     tz = LTZ._get_tz(datetime.min)
     assert tz is not None
-    assert tz.zone == 'America/New_York'
+
+
+def test_policies() -> None:
+    getzone = lambda dt: getattr(dt.tzinfo, 'zone')
+
+    naive = D('20170730 10:00:00')
+    # actual timezone at the time
+    assert getzone(TZ.localize(naive)) == 'Europe/Rome'
+
+    z = pytz.timezone('America/New_York')
+    aware = z.localize(naive)
+
+    assert getzone(TZ.localize(aware)) == 'America/New_York'
+
+    assert getzone(TZ.localize(aware, policy='convert')) == 'Europe/Rome'
+
+
+    with pytest.raises(RuntimeError):
+        assert TZ.localize(aware, policy='throw')
 
 
 def D(dstr: str) -> datetime:
@@ -91,4 +107,9 @@ def prepare(tmp_path: Path):
         # note: order doesn't matter, will be sorted in the data provider
     config.location = location # type: ignore
 
+    class time:
+        class tz:
+            pass # just rely on the default..
+    config.time = time # type: ignore
+
     yield