add_safe_accessors_decorator.py

# License: MIT
from collections.abc import Callable
from functools import wraps
from typing import Any, TypeVar

from django.db import models
from django.db.models.fields import reverse_related

T = TypeVar("T", bound=models.Model)


def add_safe_accessors(cls: type[T]) -> type[T]:
    """
    This class decorator enriches models with OneToOneField relationships by adding a 'safe_<fieldname>' attribute.
    Normally, accessing a OneToOneField incurs a database query,
    potentially raising a DoesNotExist exception if no related object is found.
    The 'safe_<fieldname>' attribute mitigates this by safely accessing the field,
    returning None instead of throwing an exception when the related object is absent.

    Recommend to define fields like this:
    ```
    class User(models.Model):
        ...
        safe_profile_image: models.OneToOneField[ProfileImage, Any] | None
    ```
    """
    original_init = cls.__init__

    @wraps(original_init)
    def __init__(self, *args, **kwargs):
        original_init(self, *args, **kwargs)
        for field in self._meta.get_fields():
            if isinstance(field, reverse_related.OneToOneRel):
                safe_field_name = f"safe_{field.name}"

                def make_safe_accessor(field_name: str) -> Callable[[T], Any]:
                    def safe_accessor(self: T) -> Any:
                        return getattr(self, field_name, None)

                    return safe_accessor

                setattr(cls, safe_field_name, property(make_safe_accessor(field.name)))

    setattr(cls, "__init__", __init__)
    return cls