Merge pull request #105 from Point72/pit/datetime

Add DatetimeContext and allow datetimes to be passed to DateContext

Author: timkpaine

ccflow/context.py

@@ -13,6 +13,7 @@
     "NullContext",
     "GenericContext",
     "DateContext",
+    "DatetimeContext",
     "EntryTimeContext",
     "DateRangeContext",
     "VersionedDateContext",
@@ -96,6 +97,19 @@ def _date_context_validator(cls, v, handler, info):
         return handler(v)
 
 
+class DatetimeContext(ContextBase):
+    dt: datetime
+
+    @model_validator(mode="wrap")
+    def _datetime_context_validator(cls, v, handler, info):
+        if cls is DatetimeContext and not isinstance(v, (DatetimeContext, dict)):
+            if isinstance(v, (tuple, list)) and len(v) == 1:
+                v = v[0]
+
+            v = DatetimeContext(dt=v)
+        return handler(v)
+
+
 class EntryTimeContext(ContextBase):
     entry_time_cutoff: Optional[datetime] = None
 

ccflow/tests/test_context.py

@@ -1,12 +1,13 @@
-from datetime import date, timedelta
+from datetime import date, datetime, timedelta, timezone
 from unittest import TestCase
 
 import pandas as pd
-from pydantic import BaseModel, ValidationError
+from pydantic import TypeAdapter, ValidationError
 
 from ccflow.context import (
     DateContext,
     DateRangeContext,
+    DatetimeContext,
     FreqContext,
     FreqDateContext,
     FreqDateRangeContext,
@@ -24,14 +25,6 @@
 from ccflow.result import GenericResult
 
 
-class MyModel(BaseModel):
-    context: DateContext
-
-
-class MyRangeModel(BaseModel):
-    context: DateRangeContext
-
-
 class TestContexts(TestCase):
     def test_null_context(self):
         n1 = NullContext()
@@ -54,13 +47,32 @@ def test_date_validation(self):
         self.assertEqual(DateContext(date="-1d"), c1)
         self.assertRaises(ValueError, DateContext, date="foo")
 
-        # Test coercion to DateContext on nested models
-        self.assertEqual(MyModel(context={"date": date.today()}).context, c)
-        self.assertEqual(MyModel(context=date.today()).context, c)
-        self.assertEqual(MyModel(context=str(date.today())).context, c)
-        self.assertEqual(MyModel(context="0d").context, c)
-        self.assertEqual(MyModel(context="-1d").context, c1)
-        self.assertRaises(ValueError, MyModel, context="foo")
+        # Test validation from other types
+        self.assertEqual(TypeAdapter(DateContext).validate_python({"date": date.today()}), c)
+        self.assertEqual(TypeAdapter(DateContext).validate_python(date.today()), c)
+        self.assertEqual(TypeAdapter(DateContext).validate_python([date.today()]), c)
+        self.assertEqual(TypeAdapter(DateContext).validate_python(str(date.today())), c)
+        self.assertEqual(TypeAdapter(DateContext).validate_python("0d"), c)
+        self.assertEqual(TypeAdapter(DateContext).validate_python("-1d"), c1)
+        self.assertRaises(ValueError, TypeAdapter(DateContext).validate_python, "foo")
+
+        # Test validation from datetime (not normally allowed by pydantic)
+        dt = datetime(2022, 1, 1, 12, tzinfo=timezone.utc)
+        self.assertEqual(TypeAdapter(DateContext).validate_python(dt), DateContext(date=dt.date()))
+
+    def test_datetime_validation(self):
+        dt = datetime(2022, 1, 1, 12, 0, tzinfo=timezone.utc)
+        c = DatetimeContext(dt=dt)
+        self.assertEqual(DatetimeContext(dt=str(dt)), c)
+        self.assertEqual(DatetimeContext(dt=dt.date()), DatetimeContext(dt=datetime(2022, 1, 1)))
+
+        # Test validation from other types
+        self.assertEqual(TypeAdapter(DatetimeContext).validate_python({"dt": dt}), c)
+        self.assertEqual(TypeAdapter(DatetimeContext).validate_python(dt), c)
+        self.assertEqual(TypeAdapter(DatetimeContext).validate_python([dt]), c)
+        self.assertEqual(TypeAdapter(DatetimeContext).validate_python(str(dt)), c)
+        self.assertEqual(TypeAdapter(DatetimeContext).validate_python(dt.isoformat()), c)
+        self.assertRaises(ValueError, TypeAdapter(DatetimeContext).validate_python, "foo")
 
     def test_coercion(self):
         d = DateContext(date=date(2022, 1, 1))
@@ -76,10 +88,11 @@ def test_date_range(self):
         self.assertEqual(DateRangeContext(start_date="-1d", end_date="0d"), c)
         self.assertRaises(ValueError, DateRangeContext, start_date="foo", end_date=d1)
 
-        # Test coercion to DateContext on nested models
-        self.assertEqual(MyRangeModel(context={"start_date": d0, "end_date": d1}).context, c)
-        self.assertEqual(MyRangeModel(context=("-1d", "0d")).context, c)
-        self.assertEqual(MyRangeModel(context=["-1d", "0d"]).context, c)
+        # Test validation from other types
+        self.assertEqual(TypeAdapter(DateRangeContext).validate_python({"start_date": d0, "end_date": d1}), c)
+        self.assertEqual(TypeAdapter(DateRangeContext).validate_python(("-1d", "0d")), c)
+        self.assertEqual(TypeAdapter(DateRangeContext).validate_python(["-1d", "0d"]), c)
+        self.assertEqual(TypeAdapter(DateRangeContext).validate_python(["-1d", datetime.now()]), c)
 
     def test_freq(self):
         self.assertEqual(

ccflow/validators.py

@@ -1,7 +1,7 @@
 """This module contains common validators."""
 
 import logging
-from datetime import date
+from datetime import date, datetime
 from typing import Any, Dict, Optional
 
 import pandas as pd
@@ -11,13 +11,15 @@
 
 
 def normalize_date(v: Any) -> Any:
-    """Validator that will convert string offsets to date based on today."""
+    """Validator that will convert string offsets to date based on today, and convert datetime to date."""
     if isinstance(v, str):  # Check case where it's an offset
         try:
             timestamp = pd.tseries.frequencies.to_offset(v) + date.today()
             return timestamp.date()
         except ValueError:
             pass
+    if isinstance(v, datetime):
+        return v.date()
     return v