Marshall class¶

def model_dump(self: BaseModel, show_pk: bool | None = None, **kwargs: Any) -> dict[str, Any]:
    """
    An updated and enhanced version of the Pydantic `model_dump` method.

    This method provides fine-grained control over how the model's data is
    serialized into a dictionary. It specifically addresses:
    -   Enforcing the inclusion of primary key fields.
    -   Correctly handling fields marked for exclusion.
    -   Applying custom logic for fields that retrieve their values via getters
        or require special serialization (e.g., related models, composite fields).

    Args:
        self: The instance of the Pydantic `BaseModel` on which `model_dump` is called.
        show_pk: An optional boolean flag. If `True`, the primary key field(s) will
                 always be included in the dumped dictionary, even if they are
                 otherwise excluded or not explicitly included. If `None`, the
                 default behavior defined by `self.__show_pk__` will be used.
        **kwargs: Arbitrary keyword arguments that are passed directly to the
                  underlying Pydantic `super().model_dump` method. These can
                  include `exclude`, `include`, `mode`, etc.

    Returns:
        A `dict` representing the serialized model data, with applied
        inclusions, exclusions, and special field handling.
    """
    meta = self.meta
    # Retrieve the 'exclude' argument, defaulting to None if not provided.
    # This argument can be a set of field names, a dictionary mapping field
    # names to boolean exclusion flags, or None.
    _exclude: set[str] | dict[str, Any] | None = kwargs.pop("exclude", None)

    # Initialize variables for managing field exclusions.
    # `initial_full_field_exclude` will contain names of fields to be fully excluded.
    # `exclude_passed` is a dictionary used for the first pass of `super().model_dump`.
    # `exclude_second_pass` is used for fields processed in the second pass (e.g., getters).
    if _exclude is None:
        initial_full_field_exclude: set[str] = _empty
        # Must be a writable dictionary to allow modifications.
        exclude_passed: dict[str, Any] = {}
        exclude_second_pass: dict[str, Any] = {}
    elif isinstance(_exclude, dict):
        # If `_exclude` is a dictionary, extract fields marked for full exclusion (value is True).
        initial_full_field_exclude = {k for k, v in _exclude.items() if v is True}
        # Create a copy for `exclude_passed` to avoid modifying the original `_exclude`.
        exclude_passed = copy.copy(_exclude)
        exclude_second_pass = _exclude
    else:
        # If `_exclude` is a set or list, convert it to a set for consistency.
        initial_full_field_exclude = set(_exclude)
        # Create dictionaries where all initially excluded fields are marked `True`.
        exclude_passed = dict.fromkeys(initial_full_field_exclude, True)
        exclude_second_pass = exclude_passed.copy()

    # `need_second_pass` will store field names that require a second processing pass.
    # These are typically fields with custom getters or foreign keys to models
    # that need special serialization.
    need_second_pass: set[str] = set()

    # Process fields that have special getter methods defined.
    for field_name in meta.special_getter_fields:
        # Temporarily exclude these fields from the initial `model_dump` pass.
        exclude_passed[field_name] = True
        # If a special getter field was not explicitly excluded and is not marked
        # for exclusion in its `MetaInfo`, add it to `need_second_pass`.
        if (
            field_name not in initial_full_field_exclude
            and not meta.fields[field_name].exclude
        ):
            need_second_pass.add(field_name)

    # Process foreign key fields.
    for field_name in meta.foreign_key_fields:
        field = meta.fields[field_name]
        # If the foreign key field is already fully excluded, skip further processing.
        if field_name in initial_full_field_exclude or field.exclude:
            continue
        # If the target model of the foreign key needs special serialization,
        # temporarily exclude it from the first pass and add to `need_second_pass`.
        if field.target.meta.needs_special_serialization:
            exclude_passed[field_name] = True
            need_second_pass.add(field_name)

    # Retrieve the 'include' argument, defaulting to None.
    include: set[str] | dict[str, Any] | None = kwargs.pop("include", None)
    # Determine the serialization mode ('json' or 'python').
    mode: Literal["json", "python"] = kwargs.pop("mode", "python")

    # Determine if the primary key should be shown, preferring the `show_pk`
    # argument over the model's internal `__show_pk__` attribute.
    should_show_pk = self.__show_pk__ if show_pk is None else show_pk

    # Perform the initial `model_dump` using Pydantic's default implementation.
    # This will handle most fields and apply the initial exclusion rules.

    result_dict: dict[str, Any] = super().model_dump(
        exclude=exclude_passed, include=include, mode=mode, **kwargs
    )

    # Set a context variable to control the behavior of `getattr` during the
    # second pass, often used to prevent recursive database queries.
    token = MODEL_GETATTR_BEHAVIOR.set("passdown")
    try:
        # Process fields identified in `need_second_pass`.
        for field_name in need_second_pass:
            # Skip the field if it's not a primary key (and `show_pk` is true)
            # or if it's not explicitly included (and `include` is not None).
            if not (
                (should_show_pk and field_name in self.pknames)
                or include is None
                or field_name in include
            ):
                continue

            field = meta.fields[field_name]
            try:
                # Attempt to get the value of the field, which might trigger a getter.
                retval = getattr(self, field_name)
            except AttributeError:
                # If the attribute doesn't exist (e.g., not loaded), skip it.
                continue

            sub_include = None
            # If `include` is a dictionary, extract specific inclusion rules for the sub-field.
            if isinstance(include, dict):
                sub_include = include.get(field_name, None)
                # If the sub-field is explicitly included with `True`, treat it as no specific
                # sub-inclusion (i.e., include all sub-fields by default).
                if sub_include is True:
                    sub_include = None

            # Get specific exclusion rules for the sub-field.
            sub_exclude = exclude_second_pass.get(field_name, None)
            # Ensure that a field marked for full exclusion in the first pass is not
            # unexpectedly processed in the second pass.
            assert sub_exclude is not True, "field should have been excluded"

            # If the retrieved value is another `BaseModel` (e.g., a related object),
            # recursively call `model_dump` on it.
            if isinstance(retval, BaseModel):
                retval = retval.model_dump(
                    include=sub_include, exclude=sub_exclude, mode=mode, **kwargs
                )
            else:
                # For non-BaseModel values, `sub_include` and `sub_exclude` should not be present
                # as they are only applicable to nested Pydantic models.
                assert sub_include is None, "sub include filters for no pydantic model"
                assert sub_exclude is None, "sub exclude filters for no pydantic model"
                # If the mode is 'json' and the field is not marked for unsafe JSON serialization,
                # skip it. This prevents non-serializable types from breaking JSON output.
                if mode == "json" and not getattr(field, "unsafe_json_serialization", False):
                    # skip field if it isn't a BaseModel and the mode is json and
                    # unsafe_json_serialization is not set
                    # Currently, `unsafe_json_serialization` exists only on `CompositeFields`.
                    continue

            # Determine the alias for the field in the dumped dictionary.
            # Prioritize `serialization_alias`, then `alias`, otherwise use the field name.
            alias: str = field_name
            if getattr(field, "serialization_alias", None):
                alias = cast(str, field.serialization_alias)
            elif getattr(field, "alias", None):
                alias = field.alias
            # Add the processed field and its value to the `result_dict`.
            result_dict[alias] = retval
    finally:
        # Reset the `MODEL_GETATTR_BEHAVIOR` context variable to its previous state.
        MODEL_GETATTR_BEHAVIOR.reset(token)
    return result_dict

def model_dump_json(self: BaseModel, **kwargs: Any) -> str:
    """
    Dumps the model data into a JSON string.

    This method leverages `model_dump` with `mode="json"` and then uses `orjson`
    for efficient JSON serialization, which is faster than Python's built-in `json` module.

    Args:
        self: The instance of the `BaseModel` to be dumped.
        **kwargs: Arbitrary keyword arguments passed to `model_dump`. These can
                  control inclusions, exclusions, and other dumping behaviors.

    Returns:
        A `str` representing the JSON serialization of the model data.
    """
    return orjson.dumps(self.model_dump(mode="json", **kwargs)).decode()

def _setup(self) -> Model:
    """
    Assembles a new Edgy model instance based on the marshall's current field values.
    This method is called when `self.instance` is accessed for the first time
    and no `instance` was provided during initialization.

    Returns:
        Model: The assembled Edgy model instance.

    Raises:
        RuntimeError: If the marshall is declared with `__incomplete_fields__`
                      and those fields are not populated for creating a new instance.
    """
    klass = type(self)
    if klass.__incomplete_fields__:
        # Prevent creating an instance if required fields are missing.
        raise RuntimeError(
            f"'{klass.__name__}' is an incomplete Marshall. "
            f"For creating new instances, it lacks following fields: [{', '.join(klass.__incomplete_fields__)}]."
        )
    model = cast(
        "Model", self.marshall_config["model"]
    )  # Get the associated Edgy model class.
    column = model.table.autoincrement_column  # Get the autoincrement PK column.
    exclude: set[str] = {*excludes_marshall}  # Start with default excluded fields.
    if column is not None:
        exclude.add(column.key)  # Exclude PK if it's autoincrement.

    # Dump marshall data to be used for the model instance.
    # Include only fields present in the model's `meta.fields` and not in `exclude`.
    data = self.model_dump(
        include=set(model.meta.fields.keys()).difference(exclude),
    )

    # Remove callable defaults that might have leaked from model fields.
    # Marshalls typically don't handle callable defaults directly.
    for k in list(data.keys()):
        if callable(data[k]):
            data.pop(k)

    # Pass internal flags to the model constructor.
    data["__show_pk__"] = self.__show_pk__
    data["__drop_extra_kwargs__"] = True
    return self.marshall_config["model"](**data)  # type: ignore

async def _resolve_async(self, name: str, awaitable: Awaitable) -> None:
    """
    Asynchronously resolves an awaitable and sets the result to a marshall field.

    Args:
        name (str): The name of the marshall field to set.
        awaitable (Awaitable): The awaitable object to resolve.
    """
    setattr(self, name, await awaitable)

def _resolve_serializer(self, instance: Model) -> Self:
    """
    Resolves the marshall's custom fields by populating them with data
    from the provided Edgy model instance. Handles both direct attribute
    access and method-based field resolution, including asynchronous getters.

    Args:
        instance (Model): The Edgy model instance to extract data from.

    Returns:
        BaseMarshall: The marshall instance with its fields populated.
    """
    async_resolvers = []
    # Iterate over custom fields defined in the marshall.
    for name, field in self.__custom_fields__.items():
        if not field.__is_method__:
            # For regular fields (not method-based).
            if name in instance.pknames:
                # Special handling for primary key attributes.
                attribute = getattr(instance, name)
            else:
                # Get attribute from instance, using 'source' if specified in BaseMarshallField.
                attribute = getattr(instance, field.source or name)

            if callable(attribute):
                # If the attribute is callable, execute it.
                value = attribute()
                if inspect.isawaitable(value):
                    # If the result is awaitable, add it to async_resolvers.
                    async_resolvers.append(self._resolve_async(name, value))
                    continue
            else:
                value = attribute
            setattr(self, name, value)  # Set the field value.
        elif field.__is_method__:
            # For method-based fields.
            value = self._get_method_value(name, instance)
            if inspect.isawaitable(value):
                # If the result of the method is awaitable, add to async_resolvers.
                async_resolvers.append(self._resolve_async(name, value))
                continue
            setattr(self, name, value)  # Set the field value.

    if async_resolvers:
        run_sync(run_concurrently(async_resolvers))  # Run all async resolvers synchronously.
    return self

def _get_method_value(self, name: str, instance: Model) -> Any:
    """
    Retrieves the value for a method-based marshall field. It expects a
    method named `get_<field_name>` on the marshall instance.

    Args:
        name (str): The name of the method-based field (e.g., 'full_name' for `get_full_name`).
        instance (Model): The Edgy model instance to pass to the getter method.

    Returns:
        Any: The value returned by the getter method.
    """
    func_name: str = f"get_{name}"
    func = getattr(self, func_name)  # Get the getter method from the marshall.
    return func(instance)  # Call the getter method with the model instance.

def _handle_primary_key(self, instance: Model) -> None:
    """
    Synchronizes the primary key fields of the marshall with those of the
    provided model instance after a save operation. This is crucial for
    newly created instances where the PK might be generated by the database.

    Args:
        instance (Model): The Edgy model instance from which to copy primary key values.
    """
    data = self.model_dump(
        include=set(instance.pknames)
    )  # Dump current PK values from marshall.

    if data:
        # Iterate over the primary key attribute names of the instance.
        for pk_attribute in instance.pknames:
            # Bypass `__setattr__` method to directly set the attribute value,
            # avoiding any potential Pydantic validation or custom logic for the PK.
            object.__setattr__(self, pk_attribute, getattr(instance, pk_attribute))

async def save(self) -> Self:
    """
    Calls the `save` method of the `BaseMarshall`, ensuring type hinting for `Self`.
    """
    return cast("Self", await super().save())