Lazy loading attributes with python attrs

attrs_lazy_loading.py

"""
Example for creating classes with attrs that contain lazy-loaded fields.
"""

import attr


_lazy_fields_container_name = "_lazy_fields"
_lazy_field_metadata_key = "lazy"
_lazy_loader_function_prefix = "lazy_loader__"


def lazy_attrib(**kwargs):
    """Extends attr.ib by adding metadata for handling lazy-loading.
    Can be used the same as `attr.ib`, with the same arguments.
    """
    metadata = kwargs.pop("metadata", dict())
    metadata[_lazy_field_metadata_key] = True
    return attr.ib(**kwargs, metadata=metadata)


def _lazy_getattribute(self, name):
    """Special __getattribute__ implementation that checks if the attribute being accessed
    is a lazy-loaded field, and if so, takes care of the lazy-loading by invoking the lazy-loader method
    """
    # get the lazy field container from the instance
    lazy_field_container = object.__getattribute__(self, _lazy_fields_container_name)
    if name in lazy_field_container:
        # check if lazy loader has already been invoked once
        if not lazy_field_container[name]:
            # lazy loader not yet invoked, get the lazy loader method
            lazy_loader_fuction_name = f"{_lazy_loader_function_prefix}{name}"
            try:
                lazy_loader_function = object.__getattribute__(self, lazy_loader_fuction_name)
            except AttributeError:
                raise AttributeError(
                    f"No lazy loader method found with name '{lazy_loader_fuction_name}' "
                    f"for class '{object.__getattribute__(self, '__class__').__name__}'"
                ) from None
            else:
                # invoke lazy loader
                lazy_loader_function()
                # indicate that it has been invoked
                lazy_field_container[name] = True
        else:
            # (lazy loader already invoked)
            pass
    return object.__getattribute__(self, name)


def create_lazy_fields(cls, fields):
    """Adds a new field for containing lazy attributes:
    A mapping between attribute names and a boolean flag that indicates
    whether the lazy loader callable has been invoked.
    Also, adds a new __getattribute__ method for the class,
    that takes care of the lazy loading.
    """
    new_fields = [f for f in fields]
    # noinspection PyArgumentList
    new_fields.append(attr.Attribute(
        name=_lazy_fields_container_name,
        default=attr.Factory(
            factory=lambda instance: {
                field.name: False
                for field in object.__getattribute__(instance, "__attrs_attrs__")
                if field.metadata.get(_lazy_field_metadata_key) is not None
            },
            takes_self=True
        ),
        validator=None,
        repr=False,
        cmp=None,
        eq=True,
        order=False,
        hash=True,
        init=True,
        inherited=False,
    ))
    setattr(cls, "__getattribute__", _lazy_getattribute)
    return new_fields


lazy_loader_invoke_count = 0


@attr.s(field_transformer=create_lazy_fields)
class Example:
    """This attrs class is extended with the `create_lazy_fields` field_transformer, that
    takes care of handling lazy attributes
    """

    # this is a noral field
    normal_field = attr.ib()
    # this field is to be lazy-loaded
    lazy_field = lazy_attrib(default=0)

    def lazy_loader__lazy_field(self):
        """This method must be implemented with this name so that lazy loading is performed."""
        global lazy_loader_invoke_count
        self.lazy_field = 42
        lazy_loader_invoke_count += 1


inst1 = Example(normal_field=1)
assert inst1.normal_field == 1
# check default value of lazy field
assert object.__getattribute__(inst1, "lazy_field") == 0
# lazy loading will take place at this attribute access
assert inst1.lazy_field == 42
# lazy loading has already been performed, won't be called now
assert inst1.lazy_field == 42

inst2 = Example(normal_field=2)
assert inst2.normal_field == 2
assert inst2.lazy_field == 42
assert inst2.lazy_field == 42

assert lazy_loader_invoke_count == 2