Merge pull request #134 from Point72/tkp/ctx

Add datetime contexts, add datetime validator convenience, update wiki to show other context examples

Author: timkpaine

ccflow/context.py

@@ -7,41 +7,60 @@
 
 from .base import ContextBase
 from .exttypes import Frequency
-from .validators import normalize_date
+from .validators import normalize_date, normalize_datetime
 
-__all__ = [
+__all__ = (
     "NullContext",
     "GenericContext",
     "DateContext",
     "DatetimeContext",
     "EntryTimeContext",
+    "SourceContext",
     "DateRangeContext",
+    "DatetimeRangeContext",
+    "SeededDateRangeContext",
+    "SeededDatetimeRangeContext",
     "VersionedDateContext",
+    "VersionedDatetimeContext",
     "VersionedDateRangeContext",
+    "VersionedDatetimeRangeContext",
     "FreqContext",
     "FreqDateContext",
+    "FreqDatetimeContext",
     "FreqDateRangeContext",
+    "FreqDatetimeRangeContext",
     "HorizonContext",
     "FreqHorizonContext",
     "FreqHorizonDateContext",
+    "FreqHorizonDatetimeContext",
     "FreqHorizonDateRangeContext",
-    "SeededDateRangeContext",
-    "SourceContext",
+    "FreqHorizonDatetimeRangeContext",
     "UniverseContext",
     "UniverseDateContext",
+    "UniverseDatetimeContext",
     "UniverseDateRangeContext",
+    "UniverseDatetimeRangeContext",
     "UniverseFrequencyDateRangeContext",
+    "UniverseFrequencyDatetimeRangeContext",
     "UniverseFrequencyHorizonDateRangeContext",
+    "UniverseFrequencyHorizonDatetimeRangeContext",
     "VersionedUniverseDateContext",
+    "VersionedUniverseDatetimeContext",
     "VersionedUniverseDateRangeContext",
+    "VersionedUniverseDatetimeRangeContext",
     "ModelContext",
     "ModelDateContext",
+    "ModelDatetimeContext",
     "ModelDateRangeContext",
+    "ModelDatetimeRangeContext",
     "ModelDateRangeSourceContext",
     "ModelFreqDateRangeContext",
+    "ModelFreqDatetimeRangeContext",
     "VersionedModelDateContext",
+    "VersionedModelDatetimeContext",
     "VersionedModelDateRangeContext",
-]
+    "VersionedModelDatetimeRangeContext",
+)
 
 _SEPARATOR = ","
 
@@ -100,6 +119,9 @@ def _date_context_validator(cls, v, handler, info):
 class DatetimeContext(ContextBase):
     dt: datetime
 
+    # validators
+    _normalize_dt = field_validator("dt", mode="before")(normalize_datetime)
+
     @model_validator(mode="wrap")
     def _datetime_context_validator(cls, v, handler, info):
         if cls is DatetimeContext and not isinstance(v, (DatetimeContext, dict)):
@@ -126,18 +148,38 @@ class DateRangeContext(ContextBase):
     _normalize_end = field_validator("end_date", mode="before")(normalize_date)
 
 
+class DatetimeRangeContext(ContextBase):
+    start_datetime: datetime
+    end_datetime: datetime
+
+    _normalize_start = field_validator("start_datetime", mode="before")(normalize_datetime)
+    _normalize_end = field_validator("end_datetime", mode="before")(normalize_datetime)
+
+
 class SeededDateRangeContext(DateRangeContext):
     seed: int = 1234
 
 
+class SeededDatetimeRangeContext(DatetimeRangeContext):
+    seed: int = 1234
+
+
 class VersionedDateContext(DateContext, EntryTimeContext):
     pass
 
 
+class VersionedDatetimeContext(DatetimeContext, EntryTimeContext):
+    pass
+
+
 class VersionedDateRangeContext(DateRangeContext, EntryTimeContext):
     pass
 
 
+class VersionedDatetimeRangeContext(DatetimeRangeContext, EntryTimeContext):
+    pass
+
+
 class FreqContext(ContextBase):
     freq: Frequency
 
@@ -146,10 +188,18 @@ class FreqDateContext(DateContext, FreqContext):
     pass
 
 
+class FreqDatetimeContext(DatetimeContext, FreqContext):
+    pass
+
+
 class FreqDateRangeContext(DateRangeContext, FreqContext):
     pass
 
 
+class FreqDatetimeRangeContext(DatetimeRangeContext, FreqContext):
+    pass
+
+
 class HorizonContext(ContextBase):
     horizon: Frequency
 
@@ -162,10 +212,18 @@ class FreqHorizonDateContext(DateContext, HorizonContext, FreqContext):
     pass
 
 
+class FreqHorizonDatetimeContext(DatetimeContext, HorizonContext, FreqContext):
+    pass
+
+
 class FreqHorizonDateRangeContext(DateRangeContext, HorizonContext, FreqContext):
     pass
 
 
+class FreqHorizonDatetimeRangeContext(DatetimeRangeContext, HorizonContext, FreqContext):
+    pass
+
+
 class UniverseContext(ContextBase):
     universe: str
 
@@ -174,26 +232,50 @@ class UniverseDateContext(DateContext, UniverseContext):
     pass
 
 
+class UniverseDatetimeContext(DatetimeContext, UniverseContext):
+    pass
+
+
 class UniverseDateRangeContext(DateRangeContext, UniverseContext):
     pass
 
 
+class UniverseDatetimeRangeContext(DatetimeRangeContext, UniverseContext):
+    pass
+
+
 class UniverseFrequencyDateRangeContext(DateRangeContext, FreqContext, UniverseContext):
     pass
 
 
+class UniverseFrequencyDatetimeRangeContext(DatetimeRangeContext, FreqContext, UniverseContext):
+    pass
+
+
 class UniverseFrequencyHorizonDateRangeContext(DateRangeContext, HorizonContext, FreqContext, UniverseContext):
     pass
 
 
+class UniverseFrequencyHorizonDatetimeRangeContext(DatetimeRangeContext, HorizonContext, FreqContext, UniverseContext):
+    pass
+
+
 class VersionedUniverseDateContext(VersionedDateContext, UniverseContext):
     pass
 
 
+class VersionedUniverseDatetimeContext(VersionedDatetimeContext, UniverseContext):
+    pass
+
+
 class VersionedUniverseDateRangeContext(VersionedDateRangeContext, UniverseContext):
     pass
 
 
+class VersionedUniverseDatetimeRangeContext(VersionedDatetimeRangeContext, UniverseContext):
+    pass
+
+
 class ModelContext(ContextBase):
     model: str
 
@@ -202,10 +284,18 @@ class ModelDateContext(DateContext, ModelContext):
     pass
 
 
+class ModelDatetimeContext(DatetimeContext, ModelContext):
+    pass
+
+
 class ModelDateRangeContext(DateRangeContext, ModelContext):
     pass
 
 
+class ModelDatetimeRangeContext(DatetimeRangeContext, ModelContext):
+    pass
+
+
 class ModelDateRangeSourceContext(SourceContext, ModelDateRangeContext):
     pass
 
@@ -214,9 +304,21 @@ class ModelFreqDateRangeContext(FreqDateRangeContext, ModelContext):
     pass
 
 
+class ModelFreqDatetimeRangeContext(FreqDatetimeRangeContext, ModelContext):
+    pass
+
+
 class VersionedModelDateContext(VersionedDateContext, ModelContext):
     pass
 
 
+class VersionedModelDatetimeContext(VersionedDatetimeContext, ModelContext):
+    pass
+
+
 class VersionedModelDateRangeContext(VersionedDateRangeContext, ModelContext):
     pass
+
+
+class VersionedModelDatetimeRangeContext(VersionedDatetimeRangeContext, ModelContext):
+    pass

ccflow/tests/test_validators.py

@@ -1,8 +1,9 @@
 import logging
 from datetime import date, datetime, timedelta
 from unittest import TestCase
+from zoneinfo import ZoneInfo
 
-from ccflow.validators import eval_or_load_object, load_object, normalize_date, str_to_log_level
+from ccflow.validators import eval_or_load_object, load_object, normalize_date, normalize_datetime, str_to_log_level
 
 
 class A:
@@ -23,6 +24,34 @@ def test_normalize_date(self):
         self.assertEqual(normalize_date("foo"), "foo")
         self.assertEqual(normalize_date(None), None)
 
+    def test_normalize_datetime(self):
+        today = datetime.today()
+        now = datetime.now()
+        c = datetime(today.year, today.month, today.day)
+
+        self.assertEqual(normalize_datetime(c), c)
+        self.assertEqual(normalize_datetime("0d"), c)
+
+        c1 = c - timedelta(1)
+        self.assertEqual(normalize_datetime("-1d"), c1)
+
+        self.assertEqual(normalize_datetime(now), now)
+
+        # check passthrough validation error
+        self.assertEqual(normalize_datetime("foo"), "foo")
+        self.assertEqual(normalize_datetime(None), None)
+
+        # check dict
+        self.assertEqual(
+            normalize_datetime({"dt": now.isoformat(), "tz": "US/Hawaii"}),
+            now.astimezone(tz=ZoneInfo("US/Hawaii")),
+        )
+        # check list
+        self.assertEqual(
+            normalize_datetime([now.isoformat(), "US/Hawaii"]),
+            now.astimezone(tz=ZoneInfo("US/Hawaii")),
+        )
+
     def test_load_object(self):
         self.assertEqual(load_object("ccflow.tests.test_validators.A"), A)
         self.assertIsNone(load_object(None))

ccflow/validators.py

@@ -3,6 +3,7 @@
 import logging
 from datetime import date, datetime
 from typing import Any, Dict, Optional
+from zoneinfo import ZoneInfo
 
 import pandas as pd
 from pydantic import TypeAdapter, ValidationError
@@ -11,6 +12,14 @@
 
 _DatetimeAdapter = TypeAdapter(datetime)
 
+__all__ = (
+    "normalize_date",
+    "normalize_datetime",
+    "load_object",
+    "eval_or_load_object",
+    "str_to_log_level",
+)
+
 
 def normalize_date(v: Any) -> Any:
     """Validator that will convert string offsets to date based on today, and convert datetime to date."""
@@ -31,6 +40,34 @@ def normalize_date(v: Any) -> Any:
     return v
 
 
+def normalize_datetime(v: Any) -> Any:
+    """Validator that will convert string offsets to datetime based on today, and convert datetime to date."""
+    if isinstance(v, str):  # Check case where it's an offset
+        try:
+            return (pd.tseries.frequencies.to_offset(v) + date.today()).to_pydatetime()
+        except ValueError:
+            pass
+    if isinstance(v, dict):
+        # e.g. DatetimeContext object, {"dt": datetime(...)}
+        dt = list(v.values())[0]
+        tz = list(v.values())[1] if len(v) > 1 else None
+    elif isinstance(v, list):
+        dt = v[0]
+        tz = v[1] if len(v) > 1 else None
+    else:
+        dt = v
+        tz = None
+    try:
+        dt = TypeAdapter(datetime).validate_python(dt)
+        if tz and isinstance(tz, str):
+            tz = ZoneInfo(tz)
+        if tz:
+            dt = dt.astimezone(tz)
+        return dt
+    except ValidationError:
+        return v
+
+
 def load_object(v: Any) -> Any:
     """Validator that loads an object from path if a string is provided"""
     if isinstance(v, str):