"""Module for managing and serializing Pydantic-based models with custom support. This module introduces the `PydanticBaseModel` class, which extends Pydantic’s `BaseModel` to facilitate custom serialization and deserialization for `pendulum.DateTime` objects. The main features include automatic handling of `pendulum.DateTime` fields, custom serialization to ISO 8601 format, and utility methods for converting model instances to and from dictionary and JSON formats. Key Classes: - PendulumDateTime: A custom type adapter that provides serialization and deserialization functionality for `pendulum.DateTime` objects, converting them to ISO 8601 strings and back. - PydanticBaseModel: A base model class for handling prediction records or configuration data with automatic Pendulum DateTime handling and additional methods for JSON and dictionary conversion. Classes: PendulumDateTime(TypeAdapter[pendulum.DateTime]): Type adapter for `pendulum.DateTime` fields with ISO 8601 serialization. Includes: - serialize: Converts `pendulum.DateTime` instances to ISO 8601 string. - deserialize: Converts ISO 8601 strings to `pendulum.DateTime` instances. - is_iso8601: Validates if a string matches the ISO 8601 date format. PydanticBaseModel(BaseModel): Extends `pydantic.BaseModel` to handle `pendulum.DateTime` fields and adds convenience methods for dictionary and JSON serialization. Key methods: - model_dump: Dumps the model, converting `pendulum.DateTime` fields to ISO 8601. - model_construct: Constructs a model instance with automatic deserialization of `pendulum.DateTime` fields from ISO 8601. - to_dict: Serializes the model instance to a dictionary. - from_dict: Constructs a model instance from a dictionary. - to_json: Converts the model instance to a JSON string. - from_json: Creates a model instance from a JSON string. Usage Example: # Define custom settings in a model using PydanticBaseModel class PredictionCommonSettings(PydanticBaseModel): prediction_start: pendulum.DateTime = Field(...) # Serialize a model instance to a dictionary or JSON config = PredictionCommonSettings(prediction_start=pendulum.now()) config_dict = config.to_dict() config_json = config.to_json() # Deserialize from dictionary or JSON new_config = PredictionCommonSettings.from_dict(config_dict) restored_config = PredictionCommonSettings.from_json(config_json) Dependencies: - `pendulum`: Required for handling timezone-aware datetime fields. - `pydantic`: Required for model and validation functionality. Notes: - This module enables custom handling of Pendulum DateTime fields within Pydantic models, which is particularly useful for applications requiring consistent ISO 8601 datetime formatting and robust timezone-aware datetime support. """ import json import re from typing import Any, Type import pendulum from pydantic import BaseModel, ConfigDict, TypeAdapter # Custom type adapter for Pendulum DateTime fields class PendulumDateTime(TypeAdapter[pendulum.DateTime]): @classmethod def serialize(cls, value: Any) -> str: """Convert pendulum.DateTime to ISO 8601 string.""" if isinstance(value, pendulum.DateTime): return value.to_iso8601_string() raise ValueError(f"Expected pendulum.DateTime, got {type(value)}") @classmethod def deserialize(cls, value: Any) -> pendulum.DateTime: """Convert ISO 8601 string to pendulum.DateTime.""" if isinstance(value, str) and cls.is_iso8601(value): try: return pendulum.parse(value) except pendulum.parsing.exceptions.ParserError as e: raise ValueError(f"Invalid date format: {value}") from e elif isinstance(value, pendulum.DateTime): return value raise ValueError(f"Expected ISO 8601 string or pendulum.DateTime, got {type(value)}") @staticmethod def is_iso8601(value: str) -> bool: """Check if the string is a valid ISO 8601 date string.""" iso8601_pattern = ( r"^(\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}(?:\.\d{1,3})?(?:Z|[+-]\d{2}:\d{2})?)$" ) return bool(re.match(iso8601_pattern, value)) class PydanticBaseModel(BaseModel): """Base model class with automatic serialization and deserialization of `pendulum.DateTime` fields. This model serializes pendulum.DateTime objects to ISO 8601 strings and deserializes ISO 8601 strings to pendulum.DateTime objects. """ # Enable custom serialization globally in config model_config = ConfigDict( arbitrary_types_allowed=True, use_enum_values=True, validate_assignment=True, ) # Override Pydantic’s serialization for all DateTime fields def model_dump(self, *args: Any, **kwargs: Any) -> dict: """Custom dump method to handle serialization for DateTime fields.""" result = super().model_dump(*args, **kwargs) for key, value in result.items(): if isinstance(value, pendulum.DateTime): result[key] = PendulumDateTime.serialize(value) return result @classmethod def model_construct(cls, data: dict) -> "PydanticBaseModel": """Custom constructor to handle deserialization for DateTime fields.""" for key, value in data.items(): if isinstance(value, str) and PendulumDateTime.is_iso8601(value): data[key] = PendulumDateTime.deserialize(value) return super().model_construct(data) def reset_optional(self) -> "PydanticBaseModel": """Resets all optional fields in the model to None. Iterates through all model fields and sets any optional (non-required) fields to None. The modification is done in-place on the current instance. Returns: PydanticBaseModel: The current instance with all optional fields reset to None. Example: >>> settings = PydanticBaseModel(name="test", optional_field="value") >>> settings.reset_optional() >>> assert settings.optional_field is None """ for field_name, field in self.model_fields.items(): if field.is_required is False: # Check if field is optional setattr(self, field_name, None) return self def to_dict(self) -> dict: """Convert this PredictionRecord instance to a dictionary representation. Returns: dict: A dictionary where the keys are the field names of the PydanticBaseModel, and the values are the corresponding field values. """ return self.model_dump() @classmethod def from_dict(cls: Type["PydanticBaseModel"], data: dict) -> "PydanticBaseModel": """Create a PydanticBaseModel instance from a dictionary. Args: data (dict): A dictionary containing data to initialize the PydanticBaseModel. Keys should match the field names defined in the model. Returns: PydanticBaseModel: An instance of the PydanticBaseModel populated with the data. Notes: Works with derived classes by ensuring the `cls` argument is used to instantiate the object. """ return cls.model_validate(data) @classmethod def from_dict_with_reset(cls, data: dict | None = None) -> "PydanticBaseModel": """Creates a new instance with reset optional fields, then updates from dict. First creates an instance with default values, resets all optional fields to None, then updates the instance with the provided dictionary data if any. Args: data (dict | None): Dictionary containing field values to initialize the instance with. Defaults to None. Returns: PydanticBaseModel: A new instance with all optional fields initially reset to None and then updated with provided data. Example: >>> data = {"name": "test", "optional_field": "value"} >>> settings = PydanticBaseModel.from_dict_with_reset(data) >>> # All non-specified optional fields will be None """ # Create instance with model defaults instance = cls() # Reset all optional fields to None instance.reset_optional() # Update with provided data if any if data: # Use model_validate to ensure proper type conversion and validation updated_instance = instance.model_validate({**instance.model_dump(), **data}) return updated_instance return instance def to_json(self) -> str: """Convert the PydanticBaseModel instance to a JSON string. Returns: str: The JSON representation of the instance. """ return self.model_dump_json() @classmethod def from_json(cls: Type["PydanticBaseModel"], json_str: str) -> "PydanticBaseModel": """Create an instance of the PydanticBaseModel class or its subclass from a JSON string. Args: json_str (str): JSON string to parse and convert into a PydanticBaseModel instance. Returns: PydanticBaseModel: A new instance of the class, populated with data from the JSON string. Notes: Works with derived classes by ensuring the `cls` argument is used to instantiate the object. """ data = json.loads(json_str) return cls.model_validate(data)