typing …

Author: commonism

pydantic_extra_types/epoch.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import datetime
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 import pydantic_core.core_schema
 from pydantic import GetJsonSchemaHandler
@@ -13,7 +13,6 @@
 
 class _Base(datetime.datetime):
     TYPE: str = ''
-    CLS: Optional[Callable[[Any], Any]]
     SCHEMA: pydantic_core.core_schema.CoreSchema
 
     @classmethod
@@ -28,28 +27,36 @@ def __get_pydantic_json_schema__(
     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:
-            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,
             cls.SCHEMA,
-            serialization=core_schema.wrap_serializer_function_ser_schema(f, return_schema=cls.SCHEMA),
+            serialization=core_schema.wrap_serializer_function_ser_schema(cls._f, return_schema=cls.SCHEMA),
         )
 
     @classmethod
     def _validate(cls, __input_value: Any, _: Any) -> datetime.datetime:
         return EPOCH + datetime.timedelta(seconds=__input_value)
 
+    @classmethod
+    def _f(cls, value: Any, serializer: Callable[[Any], Any]) -> Any:
+        pass
+
 
 class Number(_Base):
     TYPE = 'number'
-    CLS = None
     SCHEMA = core_schema.float_schema()
 
+    @classmethod
+    def _f(cls, value: Any, serializer: Callable[[float], float]) -> float:
+        ts = value.timestamp()
+        return serializer(ts)
+
 
 class Integer(_Base):
     TYPE = 'integer'
-    CLS = int
     SCHEMA = core_schema.int_schema()
+
+    @classmethod
+    def _f(cls, value: Any, serializer: Callable[[int], int]) -> int:
+        ts = value.timestamp()
+        return serializer(int(ts))