annotated.py

from functools import wraps


def annotated(annotation_name):
    """
    Decorator for returning pre-calculated, annotated values

    This should be used with model instance methods that fetch some
    kind of related or calculated data. If the method is called on
    a single instance in isolation, we should expect the method to
    execute and return its value. However if the method is called
    on a member of a queryset where the `attr_name` was added as an
    annotation, then it should return this value.

    In this way we can retain the simplicity of just "geting" the
    data for one object while retaining the performance of bulk
    annotations in a queryset (and without adding the calculations
    to the base queryset or changing named querysets).

    It additionally aids testing as it is typically simpler to
    develop the result from a single instance, which can then be
    compared to the result from the annotated queryset.

    Args:
        annotation_name: the attribute name of the annotation

    Returns:
        The underlying value, either from the method or from the queryset

    """

    def inner_decorator(func):
        @wraps(func)
        def func_wrapper(instance, *args, **kwargs):
            # MUST check if the attribute exists, NOT if is truthy
            # because if it exists and is `None` that IS what we
            # want to return
            if hasattr(instance, annotation_name):
                return getattr(instance, annotation_name)
            else:
                return func(instance, *args, **kwargs)
        return func_wrapper
    return inner_decorator