epoch - split into number & integer

Author: commonism

pydantic_extra_types/epoch.py

@@ -1,35 +1,53 @@
 import datetime
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
+import pydantic_core.core_schema
 from pydantic import GetJsonSchemaHandler
 from pydantic.json_schema import JsonSchemaValue
 from pydantic_core import CoreSchema, core_schema
 
 EPOCH = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)
 
 
-class Epoch(datetime.datetime):
+class _Base(datetime.datetime):
+    TYPE: str = ''
+    CLS: Optional[Callable[[Any], Any]]
+    SCHEMA: pydantic_core.core_schema.CoreSchema
+
     @classmethod
     def __get_pydantic_json_schema__(
         cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
     ) -> JsonSchemaValue:
         field_schema: dict[str, Any] = {}
-        field_schema.update(type='number', format='date-time')
+        field_schema.update(type=cls.TYPE, format='date-time')
         return field_schema
 
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source: type[Any], handler: Callable[[Any], CoreSchema]
     ) -> core_schema.CoreSchema:
         def f(value: Any, serializer: Callable[[datetime.datetime], float]) -> float:
-            return serializer(value.timestamp())
+            ts = value.timestamp()
+            return serializer(cls.CLS(ts) if cls.CLS is not None else ts)
 
         return core_schema.with_info_after_validator_function(
             cls._validate,
-            core_schema.float_schema(),
-            serialization=core_schema.wrap_serializer_function_ser_schema(f, return_schema=core_schema.float_schema()),
+            cls.SCHEMA,
+            serialization=core_schema.wrap_serializer_function_ser_schema(f, return_schema=cls.SCHEMA),
         )
 
     @classmethod
     def _validate(cls, __input_value: Any, _: Any) -> datetime.datetime:
         return EPOCH + datetime.timedelta(seconds=__input_value)
+
+
+class Number(_Base):
+    TYPE = 'number'
+    CLS = None
+    SCHEMA = core_schema.float_schema()
+
+
+class Integer(_Base):
+    TYPE = 'integer'
+    CLS = int
+    SCHEMA = core_schema.int_schema()

tests/test_epoch.py

@@ -2,15 +2,15 @@
 
 import pytest
 
-from pydantic_extra_types.epoch import Epoch
+from pydantic_extra_types import epoch
 
 
-@pytest.mark.parametrize('type_', [(int,), (float,)], ids=['integer', 'float'])
-def test_type(type_):
+@pytest.mark.parametrize('type_,cls_', [(int, epoch.Integer), (float, epoch.Number)], ids=['integer', 'number'])
+def test_type(type_, cls_):
     from pydantic import BaseModel
 
     class A(BaseModel):
-        epoch: Epoch
+        epoch: cls_
 
     now = datetime.datetime.now(tz=datetime.timezone.utc)
     ts = type_(now.timestamp())
@@ -28,11 +28,12 @@ class A(BaseModel):
     assert v['epoch'] == ts
 
 
-def test_schema():
+@pytest.mark.parametrize('cls_', [(epoch.Integer), (epoch.Number)], ids=['integer', 'number'])
+def test_schema(cls_):
     from pydantic import BaseModel
 
     class A(BaseModel):
-        dt: Epoch
+        dt: cls_
 
     v = A.model_json_schema()
-    assert (dt := v['properties']['dt'])['type'] == 'number' and dt['format'] == 'date-time'
+    assert (dt := v['properties']['dt'])['type'] == cls_.TYPE and dt['format'] == 'date-time'