Python Recursive Decorators

In Python decorators are a useful tool for changing the behaviour of functions without modifying the original function. For a recent release of openapi-SQLAlchemy, which added support for references among table columns, I used decorators to de-reference columns. I needed a way of supporting cases where a reference to a column was just a reference to another column. The solution was to essentially keep applying the decorator until the column was actually found.

What is a Column Reference?

If you are not familiar with openapi specifications, a simple schema for an object might be the following:

components:
  schemas:
    Employee:
      description: Person that works for a company.
      type: object
      properties:
        id:
          type: integer
          description: Unique identifier for the employee.
          example: 0
        name:
          type: string
          description: The name of the employee.
          example: David Andersson.

To be able to re-use the definition of the id property for another schema, you can do the following:

components:
  schemas:
    Id:
      type: integer
      description: Unique identifier for the employee.
      example: 0
    Employee:
      description: Person that works for a company.
      type: object
      properties:
        id:
          $ref: "#/components/schemas/Id"
        name:
          type: string
          description: The name of the employee.
          example: David Andersson.

In this case, the id property just references the Id schema. This could be done for other property definitions where the same schema applies to reduce duplicate schema definitions.

The Simple Case

The openapi-SQLAlchemy package allows you to map openapi schemas to SQLAlchemy models where an object becomes a table and a property becomes a column. The architecture of the package is broken into a factory for tables which then calls a column factory for each property of an object. The problem I had to solve was that ,when a reference is encountered, the column factory gets called with the following dictionary as the schema for the column (for the id column in this case):

{"$ref": "#/components/schemas/Id"}

instead of the schema for the column. Apart from doing some basic checks, what needs to happen is that the schema of Id needs to be found and the column factory be called with that schema instead of the reference. A perfect case for a decorator! The following is the code (minus some input checks):

_REF_PATTER = re.compile(r"^#\/components\/schemas\/(\w+)$")

def resolve_ref(func):
    """Resolve $ref schemas."""

    def inner(schema, schemas, **kwargs):
        """Replace function."""
        # Checking for $ref
        ref = schema.get("$ref")
        if ref is None:
            return func(schema=schema, **kwargs)

        # Retrieving new schema
        match = _REF_PATTER.match(ref)
        schema_name = match.group(1)
        ref_schema = schemas.get(schema_name)

        return func(schema=ref_schema, **kwargs)

    return inner

The first step is to check if the schema is a reference schema and call the column factory if not (lines 3-4). If it is a reference, then the referenced schema needs to be retrieved (lines 14-16) and the factory called with the referenced schema instead (line 18).

The Recursive Case

The problem with the simple decorator is that the following openapi specification is valid:

components:
  schemas:
    Id:
      $ref: "#/components/schemas/RefId"
    RefId:
      type: integer
      description: Unique identifier for the employee.
      example: 0
    Employee:
      description: Person that works for a company.
      type: object
      x-tablename: employee
      properties:
        id:
          $ref: "#/components/schemas/Id"
        name:
          type: string
          description: The name of the employee.
          example: David Andersson.

Noe that the Id schema just references the RefId schema. When this schema is used, the column factory would now be called with:

{"$ref": "#/components/schemas/RefId"}

That looks a lot like the original problem! The solution is that the decorator needs to be applied again. This also looks like a case for recursive programming. In recursive programming, you have the base case and the recursive case. The base case sounds a lot like the code that checks for whether the schema is a reference (lines 3-4 above). The recursive case needs to take a step towards the base case and apply the function again. In this case this is done by resolving one reference and then calling the decorator again. This means that on line 18 above we need to somehow apply the decorator again. The solution is actually quite simple, the code needs to be changed to the following:

_REF_PATTER = re.compile(r"^#\/components\/schemas\/(\w+)$")

def resolve_ref(func):
    """Resolve $ref schemas."""

    def inner(schema, schemas, **kwargs):
        """Replace function."""
        # Checking for $ref
        ref = schema.get("$ref")
        if ref is None:
            return func(schema=schema, **kwargs)

        # Retrieving new schema
        match = _REF_PATTER.match(ref)
        schema_name = match.group(1)
        ref_schema = schemas.get(schema_name)

        return inner(schema=ref_schema, schemas=schemas, **kwargs)

    return inner

The change is that inner is called instead of the original function func. We then also need to pass in all the required arguments for inner, which in this case, means passing through the schemas which func does not need.

Conclusion

Recursive decorators combine the ideas behind decorators and recursive programming. If the problem a decorator solves can be broken into steps where, after each step, the decorator might need to be applied again, you might have a case for a recursive decorator. The decorator must then implement a base case where the decorator is not applied again and a recursive case where one step is taken towards the base case and the decorator is applied again.

Testing the Testing Guard Decorator in Python

Testing functions in Python that have a decorator applied to them is not ideal as the tests have to be written to take into account the decorator. The testing guard decorator can help remove the problem by selectively running the function with or without the decorator by detecting the environment in which it is run. When you make the decorator a part of your project you might like to write tests for the testing guard decorator itself.

Testing Guard Decorator

As a reminder, the following is the decorator code for the testing guard.

# testing_guard.py
"""Demonstrates a guard decorator."""

import os


def testing_guard(decorator_func):
    """
    Decorator that only applies another decorator if the appropriate
    environment variable is not set.

    Args:
        decorator_func: The function that applies the decorator.

    Returns:
        Function that dynamically decides whether to apply the decorator based
        on the environment.
    """
    def replacement(original_func):
        """
        Function that is used as the decorator function instead of the
        decorator function.

        Args:
            original_func: The function being decorated.

        Returns:
            Function that dynamically picks between the original and
            decorated function depending on the environment.
        """
        # Creating decorated function
        decorated_func = decorator_func(original_func)

        def apply_guard(*args, **kwargs):
            """
            Dynamically picks between decorated and original function based on
            environment.
            """
            if os.getenv('TESTING') is not None:
                # Use original function
                return original_func(*args, **kwargs)
            # Use decorated function
            return decorated_func(*args, **kwargs)

        return apply_guard

    return replacement

The key pieces of functionality are on lines 39-43. Line 39 is the check for whether the function is being executed in the testing environment. If it is, the original function is called without the decorator on line 41. If it isn’t, then the decorated function is called on line 43. The key reason that the decorator works is that it gets access to both the decorator and original function and that it gets to run code every time the decorated function is run.

Testing the Testing Guard Decorator

There are 2 scenarios to test. The first is that the decorator is executed in the test environment, and the second is when it is executed in the non-test environment. In both cases we care about both that the correct function is called with the correct arguments (and that the other function is not called) and that the return value from the correct function is returned.

Let’s first define 2 fixtures that will help during testing. We will repeatedly need some generic *args and **kwargs and the particular value we choose for either doesn’t add anything to help understand the tests themselves so it is better to hide that detail in a fixture.

@pytest.fixture(scope='session')
def args():
    """Generates function *args"""
    return ('arg1', 'arg2')


@pytest.fixture(scope='session')
def kwargs():
    """Generates function **kwargs"""
    return {'kwargs1': 'kwarg1', 'kwarg2': 'kwarg2'}

Test Environment

When the decorator is executed in the test environment, it should call the original function and return the return value of the original function. It should not call the decorator return value. To focus the tests, let’s split those into three separate tests.

The first test just checks that the decorator return value is not called:

def test_testing_guard_set_decorated_call(monkeypatch):
    """
    GIVEN TESTING environment variable is set and mock decorator
    WHEN decorator is applied to a function after decorating it with the
        testing guard and calling the decorated function
    THEN decorator return value is not called.
    """
    # Setting TESTING environment variable
    monkeypatch.setenv('TESTING', '')
    # Defining mock decorator
    mock_decorator = mock.MagicMock()

    # Decorating with testing guard and calling
    guarded_mock_decorator = testing_guard(mock_decorator)
    # Applying decorator
    mock_decorated_func = guarded_mock_decorator(mock.MagicMock())
    # Calling function
    mock_decorated_func()

    # Checking decorator call
    mock_decorator.return_value.assert_not_called()

To understand the test, some understanding of Python decorators is required. A decorator is another function whose return value replaces the function it is decorating. A decorator could, for example, return the print function and then the function being decorator would never be called. In most cases, however, the function being decorator is called in the body of the function the decorator returns.

On line 9 of the test the testing environment is setup. On line 11 a mock function that will serve as the decorator to which the testing guard is applied is defined. On line 14 the testing guard is applied to the mock decorator. Then a mock function is decorated with guarded mock decorator on line 16 and the decorated function is called on line 18. On line 21 it is checked that the return value of the decorator, which would usually be called instead of the original function, is not called since, in the testing environment, the original function should always be called.

The next test checks that the original function is called with the correct arguments. For this we will ned the args and kwargs fixtures.

def test_testing_guard_set_func_call(monkeypatch, args, kwargs):
    """
    GIVEN TESTING environment variable is set, mock function and args and
        kwargs
    WHEN a decorator is applied to the function after decorating it with the
        testing guard and calling the decorated function with args and kwargs
    THEN function is called with args and kwargs.
    """
    # Setting TESTING environment variable
    monkeypatch.setenv('TESTING', '')
    # Defining mock decorator
    mock_func = mock.MagicMock()

    # Decorating with testing guard and calling
    mock_decorator = mock.MagicMock()
    guarded_mock_decorator = testing_guard(mock_decorator)
    # Applying decorator
    mock_decorated_func = guarded_mock_decorator(mock_func)
    # Calling function
    mock_decorated_func(*args, **kwargs)

    # Checking decorator call
    mock_func.assert_called_once_with(*args, **kwargs)

This tests is very similar to the first test but, instead of keeping track of the decorator, the original function is defined on line 12. The procedure on lines 14-20 is very similar to the first test with the only difference being that the mock decorated function is called with the args and kwargs fixtures. On line 23 the original function call is checked.

The final test checks the return value of the mock decorated function call. It is much like the previous test, except that it doesn’t pass in any args nor kwargs.

def test_testing_guard_set_return(monkeypatch):
    """
    GIVEN TESTING environment variable is set and mock function
    WHEN a decorator is applied to the function after decorating it with the
        testing guard and calling the decorated function
    THEN the return value is the function return value.
    """
    # Setting TESTING environment variable
    monkeypatch.setenv('TESTING', '')
    # Defining mock function
    mock_func = mock.MagicMock()

    # Decorating with testing guard and calling
    mock_decorator = mock.MagicMock()
    guarded_mock_decorator = testing_guard(mock_decorator)
    # Applying decorator
    mock_decorated_func = guarded_mock_decorator(mock_func)
    # Calling function
    return_value = mock_decorated_func()

    # Checking decorator call
    assert return_value == mock_func.return_value

Non-Test Environment

The second series of tests are outside the test environment. Since the tests for the testing guard are likely running inside the test environment, it is usually best to deactivate the testing environment as a part of the test as you may want to have the testing environment active by default during your tests.

The tests are very similar to the tests in the test environment. The difference is that now the decorator return value should be called, the original function should not be called and the return value of the decorator return value should be returned. The tests are shown below.

def test_testing_guard_not_set_decorated_call(monkeypatch, args, kwargs):
    """
    GIVEN TESTING environment variable is not set, mock decorator and args and
        kwargs
    WHEN decorator is applied to a function after decorating it with the
        testing guard and calling the decorated function with args and kwargs
    THEN decorator return value is called with args and kwargs.
    """
    # Removing TESTING environment variable
    monkeypatch.delenv('TESTING', raising=False)
    # Defining mock decorator
    mock_decorator = mock.MagicMock()

    # Decorating with testing guard and calling
    guarded_mock_decorator = testing_guard(mock_decorator)
    # Applying decorator
    mock_decorated_func = guarded_mock_decorator(mock.MagicMock())
    # Calling function
    mock_decorated_func(*args, **kwargs)

    # Checking decorator call
    mock_decorator.return_value.assert_called_once_with(*args, **kwargs)


def test_testing_guard_not_set_return(monkeypatch):
    """
    GIVEN TESTING environment variable is not set and mock decorator
    WHEN decorator is applied to a function after decorating it with the
        testing guard and calling the decorated function
    THEN the return value is the decorator's return value return value.
    """
    # Removing TESTING environment variable
    monkeypatch.delenv('TESTING', raising=False)
    # Defining mock decorator
    mock_decorator = mock.MagicMock()

    # Decorating with testing guard and calling
    guarded_mock_decorator = testing_guard(mock_decorator)
    # Applying decorator
    mock_decorated_func = guarded_mock_decorator(mock.MagicMock())
    # Calling function
    return_value = mock_decorated_func()

    # Checking decorator call
    assert return_value == mock_decorator.return_value.return_value


def test_testing_guard_not_set_func_call(monkeypatch):
    """
    GIVEN TESTING environment variable is not set and mock function
    WHEN a decorator is applied to the function after decorating it with the
        testing guard and calling the decorated function
    THEN function is not called.
    """
    # Removing TESTING environment variable
    monkeypatch.delenv('TESTING', raising=False)
    # Defining mock decorator
    mock_func = mock.MagicMock()

    # Decorating with testing guard and calling
    mock_decorator = mock.MagicMock()
    guarded_mock_decorator = testing_guard(mock_decorator)
    # Applying decorator
    mock_decorated_func = guarded_mock_decorator(mock_func)
    # Calling function
    mock_decorated_func()

    # Checking decorator call
    mock_func.assert_not_called()

This demonstrates how to test the testing guard. The setup and teardown of your test environment might be more complicated in which case your tests would have to reflect that.

Testing Decorated Python Functions

Decorators are a great way of adding functionality to a function with minimal impact on the function itself. On top of that, decorator logic can be re-used on other functions that also require the new functionality. For example the following decorator prints a message to standard output every time a function is called.

# plain_main.py
"""Demonstrates a simple decorator."""


def decorator(func):
    """
    A simple decorator that adds printing a message on a function call.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """
    def inner(*args, **kwargs):
        """Function that is called instead of original function."""
        print('The decorator was called.')
        return func(*args, **kwargs)

    return inner


@decorator
def main():
    print('The main function was called.')


if __name__ == '__main__':
    print('Calling the main function.')
    main()

$ python3 plain_main.py 
Calling the main function.
The decorator was called.
The main function was called.

The drawback of decorators is that the decorator is applied as soon as the interpreter reaches the function definition and it is hard to access the original function without the decorator applied. This might be desirable during testing where testing of the function and the decorator should be separated.

Adding Testing Guard Logic to a Decorator

The solution is to optionally skip the decorator logic if a certain condition is met that is only true during testing. For example, skip the decorator logic if the TESTING environment variable is set.

# check_main.py
"""Demonstrates a decorator with a testing guard."""

import os


def decorator(func):
    """
    A simple decorator that adds printing a message on a function call unless
    the TESTING environment variable is set.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """
    def inner(*args, **kwargs):
        """Function that is called instead of original function."""
        # Checking for TESTING environment variable
        if os.getenv('TESTING') is not None:
            # Skipping decortor logic
            return func(*args, **kwargs)

        # Running decorator logic
        print('The decorator was called.')
        return func(*args, **kwargs)

    return inner


@decorator
def main():
    print('The main function was called.')


if __name__ == '__main__':
    print('Calling the main function without TESTING set.')
    main()

    print('Calling the main function with TESTING set.')
    os.environ['TESTING'] = ''
    main()

$ python3 check_main.py 
Calling the main function without TESTING set.
The decorator was called.
The main function was called.
Calling the main function with TESTING set.
The main function was called.

As you can see, the decorator logic was executed under normal circumstances (main call on line 39) and was skipped when the TESTING environment variable was set (main call on line 43). The reason was because of the guard statement on line 21 that checks for the TESTING environment variable.

Guard Decorator

You might now say: “great, thank you David. Now I have to rewrite all of my decorator functions!” Ah, but you don’t. If you stay with me through a little more complex decorator code, you won’t have to! The idea is to write a decorator that modifies another decorator’s behaviour.

# guard_main.py
"""Demonstrates a guard decorator."""

import os


def testing_guard(decorator_func):
    """
    Decorator that only applies another decorator if the TESTING environment
    variable is not set.

    Args:
        decorator_func: The decorator function.

    Returns:
        Function that calls a function after applying the decorator if TESTING
        environment variable is not set and calls the plain function if it is set.
    """
    def replacement(original_func):
        """Function that is called instead of original function."""
        def apply_guard(*args, **kwargs):
            """Decides whether to use decorator on function call."""
            if os.getenv('TESTING') is not None:
                return original_func(*args, **kwargs)
            return decorator_func(original_func)(*args, **kwargs)

        return apply_guard
    return replacement


@testing_guard
def decorator(func):
    """
    A simple decorator that adds printing a message on a function call.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """
    def replacement(*args, **kwargs):
        """Function that is called instead of original function."""
        print('The decorator was called.')
        return func(*args, **kwargs)

    return replacement


@decorator
def main():
    print('The main function was called.')


if __name__ == '__main__':
    print('Calling the main function without TESTING set.')
    main()

    print('Calling the main function with TESTING set.')
    os.environ['TESTING'] = ''
    main()

$ python3 guard_main.py 
Calling the main function without TESTING set.
The decorator was called.
The main function was called.
Calling the main function with TESTING set.
The main function was called.

As you can see, the behaviour of the code is exactly the same but the decorator is in its original form. The reason this works is because the guard decorator gets to intercept each function call and can then decide whether to first apply the decorator or call the plain function on lines 23-25.

On top of not having to re-write decorator functions which you don’t want to execute during testing, you also get to separate the logic that determines whether the decorator is applied from the decorator logic which will reduce the chances of accidentally executing decorator logic as you make changes to the decorator. It also helps you write clear unit tests for both the decorator and the guard decorator.

Working with Pytest

The last consideration is how do you apply this in practice. I would argue that, unless you are testing the decorator or guard decorator, you should always have the TESTING environment variable set. This ensures that you are only testing function logic and not decorator logic. You can achieve this by putting a fixture in your root conftest.py file with autouse set to True.

@pytest.fixture(scope='function', autouse=True)
def set_testing(monkeypatch):
    """Sets the TESTING environment variable."""
    monkeypatch.setenv('TESTING', '')

When you are testing the decorator you would always ensure that the TESTING environment variable is not set. You can achieve that using a fixture that clears the TESTING environment variable that also has autouse set to True as a part of the file that tests the decorator.

@pytest.fixture(scope='function', autouse=True)
def delete_testing(monkeypatch):
    """Deletes the TESTING environment variable."""
    monkeypatch.delenv('TESTING', raising=False)

Finally, for testing the guard decorator, make whether the TESTING environment variable is set part of the tests themselves. Don’t be shy about overriding the functionality of any autouse fixtures as a part of the test function to demonstrate what the intended state of the test is clearly.

def test_guard_decorator_testing_set(monkeypatch):
    """
    GIVEN TESTING environment variable set and ...
    WHEN ...
    THEN ...
    """
    # Setting TESTING environment variable
    monkeypatch.setenv('TESTING', '')

    # Other test code


def test_guard_decorator_testing_not_set(monkeypatch):
    """
    GIVEN TESTING environment variable is not set and ...
    WHEN ...
    THEN ...
    """
    # Setting TESTING environment variable
    monkeypatch.delenv('TESTING', raising=False)

    # Other test code

Thats it! Consider whether environment variables is the best way of indicating that decorator logic should be skipped during testing and also what the best name of the environment variable is for you. I hope this was useful to you!