def add_to_repr()

in pygenie/jobs/utils.py [0:0]

• 23 lines of code
• 6 McCabe index (conditional complexity)

def add_to_repr(mode):
    """
    Decorator for adding a method call to an object's repr. The decorator
    assumes the object has a Repr object instantiated in the object's
    repr_obj attribute. If the object does not have this attribute, will do
    nothing.

    When adding a call to the repr list, the following modes are available:
        append: methods will just be added to the repr list.
        overwrite: if the method was previously called, overwrite the call in
            the repr list.

    Args:
        mode (str): The mode when adding to the object's repr list ('append',
            'overwrite', etc).
    """

    assert is_str(mode) and mode.lower() in REPR_MODES, \
        "invalid mode for adding to repr ('{}'), use: {}" \
            .format(mode, REPR_MODES)

    mode = mode.lower()

    def real_decorator(func):
        """Wrap the function call."""

        @_remove_wrapper_from_args
        def wrapper(*args, **kwargs):
            """Adds the method call to an object's repr list."""

            repr_obj = getattr(args[0], 'repr_obj', None)

            if repr_obj:
                if mode in REPR_OVERWRITE_MODES:
                    repr_obj.remove(r'{}('.format(func.__name__))
                repr_obj.append(func_name=func.__name__,
                                args=args[1:],
                                kwargs=kwargs)

            try:
                return func(*args, **kwargs)
            except Exception:
                if repr_obj:
                    repr_obj.pop()
                raise

        return decorator(wrapper, func)

    return real_decorator