|
| 1 | +""" |
| 2 | +Native Pendulum DateTime object implementation. This is a copy of the Pendulum DateTime object, but with a Pydantic |
| 3 | +CoreSchema implementation. This allows Pydantic to validate the DateTime object. |
| 4 | +""" |
| 5 | + |
| 6 | +try: |
| 7 | + from pendulum import DateTime as _DateTime |
| 8 | + from pendulum import parse |
| 9 | +except ModuleNotFoundError: # pragma: no cover |
| 10 | + raise RuntimeError( |
| 11 | + 'The `pendulum_dt` module requires "pendulum" to be installed. You can install it with "pip install pendulum".' |
| 12 | + ) |
| 13 | +from typing import Any, List, Type |
| 14 | + |
| 15 | +from pydantic import GetCoreSchemaHandler |
| 16 | +from pydantic_core import PydanticCustomError, core_schema |
| 17 | + |
| 18 | + |
| 19 | +class DateTime(_DateTime): |
| 20 | + """ |
| 21 | + A `pendulum.DateTime` object. At runtime, this type decomposes into pendulum.DateTime automatically. |
| 22 | + This type exists because Pydantic throws a fit on unknown types. |
| 23 | +
|
| 24 | + ```python |
| 25 | + from pydantic import BaseModel |
| 26 | + from pydantic_extra_types.pendulum_dt import DateTime |
| 27 | +
|
| 28 | + class test_model(BaseModel): |
| 29 | + dt: DateTime |
| 30 | +
|
| 31 | + print(test_model(dt='2021-01-01T00:00:00+00:00')) |
| 32 | +
|
| 33 | + #> test_model(dt=DateTime(2021, 1, 1, 0, 0, 0, tzinfo=FixedTimezone(0, name="+00:00"))) |
| 34 | + ``` |
| 35 | + """ |
| 36 | + |
| 37 | + __slots__: List[str] = [] |
| 38 | + |
| 39 | + @classmethod |
| 40 | + def __get_pydantic_core_schema__(cls, source: Type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: |
| 41 | + """ |
| 42 | + Return a Pydantic CoreSchema with the Datetime validation |
| 43 | +
|
| 44 | + Args: |
| 45 | + source: The source type to be converted. |
| 46 | + handler: The handler to get the CoreSchema. |
| 47 | +
|
| 48 | + Returns: |
| 49 | + A Pydantic CoreSchema with the Datetime validation. |
| 50 | + """ |
| 51 | + return core_schema.no_info_wrap_validator_function(cls._validate, core_schema.datetime_schema()) |
| 52 | + |
| 53 | + @classmethod |
| 54 | + def _validate(cls, value: Any, handler: core_schema.ValidatorFunctionWrapHandler) -> Any: |
| 55 | + """ |
| 56 | + Validate the datetime object and return it. |
| 57 | +
|
| 58 | + Args: |
| 59 | + value: The value to validate. |
| 60 | + handler: The handler to get the CoreSchema. |
| 61 | +
|
| 62 | + Returns: |
| 63 | + The validated value or raises a PydanticCustomError. |
| 64 | + """ |
| 65 | + # if we are passed an existing instance, pass it straight through. |
| 66 | + if isinstance(value, _DateTime): |
| 67 | + return handler(value) |
| 68 | + |
| 69 | + # otherwise, parse it. |
| 70 | + try: |
| 71 | + data = parse(value) |
| 72 | + except Exception as exc: |
| 73 | + raise PydanticCustomError('value_error', 'value is not a valid timestamp') from exc |
| 74 | + return handler(data) |
0 commit comments