From 28ef1037d2edb314b482dc39e7f97c376bf23d1c Mon Sep 17 00:00:00 2001 From: commonism Date: Wed, 20 Nov 2024 21:49:19 +0100 Subject: [PATCH] :sparkles: add new type: Epoch Unix Timestamp (#240) --- pydantic_extra_types/epoch.py | 95 +++++++++++++++++++++++++++++++++++ tests/test_epoch.py | 39 ++++++++++++++ tests/test_json_schema.py | 35 +++++++++++++ 3 files changed, 169 insertions(+) create mode 100644 pydantic_extra_types/epoch.py create mode 100644 tests/test_epoch.py diff --git a/pydantic_extra_types/epoch.py b/pydantic_extra_types/epoch.py new file mode 100644 index 0000000..c9dda44 --- /dev/null +++ b/pydantic_extra_types/epoch.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import datetime +from typing import Any, Callable + +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 _Base(datetime.datetime): + TYPE: str = '' + 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=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: + return core_schema.with_info_after_validator_function( + cls._validate, + 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: # pragma: no cover + raise NotImplementedError(cls) + + +class Number(_Base): + """epoch.Number parses unix timestamp as float and converts it to datetime. + + ```py + from pydantic import BaseModel + + from pydantic_extra_types import epoch + + class LogEntry(BaseModel): + timestamp: epoch.Number + + logentry = LogEntry(timestamp=1.1) + print(logentry) + #> timestamp=datetime.datetime(1970, 1, 1, 0, 0, 1, 100000, tzinfo=datetime.timezone.utc) + ``` + """ + + TYPE = 'number' + 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): + """epoch.Integer parses unix timestamp as integer and converts it to datetime. + + ``` + ```py + from pydantic import BaseModel + + from pydantic_extra_types import epoch + + class LogEntry(BaseModel): + timestamp: epoch.Integer + + logentry = LogEntry(timestamp=1) + print(logentry) + #> timestamp=datetime.datetime(1970, 1, 1, 0, 0, 1, tzinfo=datetime.timezone.utc) + ``` + """ + + TYPE = 'integer' + SCHEMA = core_schema.int_schema() + + @classmethod + def _f(cls, value: Any, serializer: Callable[[int], int]) -> int: + ts = value.timestamp() + return serializer(int(ts)) diff --git a/tests/test_epoch.py b/tests/test_epoch.py new file mode 100644 index 0000000..5cbfb62 --- /dev/null +++ b/tests/test_epoch.py @@ -0,0 +1,39 @@ +import datetime + +import pytest + +from pydantic_extra_types import epoch + + +@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: cls_ + + now = datetime.datetime.now(tz=datetime.timezone.utc) + ts = type_(now.timestamp()) + a = A.model_validate({'epoch': ts}) + v = a.model_dump() + assert v['epoch'] == ts + + b = A.model_construct(epoch=now) + + v = b.model_dump() + assert v['epoch'] == ts + + c = A.model_validate(dict(epoch=ts)) + v = c.model_dump() + assert v['epoch'] == ts + + +@pytest.mark.parametrize('cls_', [(epoch.Integer), (epoch.Number)], ids=['integer', 'number']) +def test_schema(cls_): + from pydantic import BaseModel + + class A(BaseModel): + dt: cls_ + + v = A.model_json_schema() + assert (dt := v['properties']['dt'])['type'] == cls_.TYPE and dt['format'] == 'date-time' diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py index 7929fa4..96773bc 100644 --- a/tests/test_json_schema.py +++ b/tests/test_json_schema.py @@ -11,6 +11,7 @@ from typing import Annotated import pydantic_extra_types +from pydantic_extra_types import epoch from pydantic_extra_types.color import Color from pydantic_extra_types.coordinate import Coordinate, Latitude, Longitude from pydantic_extra_types.country import CountryAlpha2, CountryAlpha3, CountryNumericCode, CountryShortName @@ -464,6 +465,40 @@ ], }, ), + ( + epoch.Integer, + { + 'title': 'Model', + 'type': 'object', + 'properties': { + 'x': { + 'title': 'X', + 'type': 'integer', + 'format': 'date-time', + }, + }, + 'required': [ + 'x', + ], + }, + ), + ( + epoch.Number, + { + 'title': 'Model', + 'type': 'object', + 'properties': { + 'x': { + 'title': 'X', + 'type': 'number', + 'format': 'date-time', + }, + }, + 'required': [ + 'x', + ], + }, + ), ], ) def test_json_schema(cls, expected):