Python Asteroids+Invaders on GitHub

We continue learning decorators using tests. Our simple decorator doesn’t handle parameters in the wrapped function. Lets figure out how to do that.

Parameters on Wrapped Function

So far, we have wrapped a simple function that takes no parameters. If our little repeat decorator is to be fully general, the functions it wraps must be allowed to have parameters. We need at least one more test. I’ll start by replicating the current repeat in the new test, and we’ll see what happens.

    def test_repeat_with_function_parameters(self):
        def repeat(times):
            def rpt(func):
                @functools.wraps(func)
                def wrapped():
                    for i in range(times):
                        func()
                return wrapped
            return rpt
        count = 0
        @repeat(3)
        def increment(by):
            nonlocal count
            count += by
        increment(2)
        assert count == 6

I suspect that this will fail saying something about too many parameters.

>       increment(2)
E       TypeError: TestDecorators.test_repeat_with_function_parameters.<locals>.increment() 
takes 0 positional arguments but 1 was given

Since it’s referring to increment() and we used the wraps feature, I reckon it’s talking about the wrapper(). Since we don’t know how many parameters we’ll get, we’ll start with this:

        def repeat(times):
            def rpt(func):
                @functools.wraps(func)
                def wrapped(*args):
                    for i in range(times):
                        func(*args)
                return wrapped
            return rpt

And our test passes. But what about keyword arguments? I rarely use them, but they are part of the language, so we need to cater for them.

The hard part is coming up with an example. Here we go, a multiplier:

    def test_repeat_with_function_parameters(self):
        def repeat(times):
            def rpt(func):
                @functools.wraps(func)
                def wrapped(*args):
                    for i in range(times):
                        func(*args)
                return wrapped
            return rpt
        count = 0
        @repeat(3)
        def increment(by, multiplier=1):
            nonlocal count
            count += by*multiplier
        increment(2)
        assert count == 6
        count = 0
        increment(2, multiplier=2)
        assert count == 12

Test fails, saying:

>       increment(2, multiplier=2)
E       TypeError: TestDecorators.test_repeat_with_function_parameters.<locals>.increment()
got an unexpected keyword argument 'multiplier'

Just as one might expect. How do we do general kw args? Like this:

        def repeat(times):
            def rpt(func):
                @functools.wraps(func)
                def wrapped(*args, **kwargs):
                    for i in range(times):
                        func(*args, **kwargs)
                return wrapped
            return rpt

We are green. Commit: test for and use args and kwargs.

Let’s reflect.

Reflection

When wrapping functions, we will generally not know what parameters the wrapped function will take. In Python, we can define a function to accept any number of standard parameters and any number of keyword parameters with:

def func(*args, **kwargs)

So our “recipe” for function decorators should generally accept *args and **kwargs and pass them to the wrapped function. That will handle all the cases from no parameter through multiple parameters of both kinds. (We could write another test with multiples, but I feel no need for it. Feel free to write one yourself, and if it fails, let me know.)

What now? We have not wrapped a class with a function-based decorator. I guess we should do that.

Wrap Class with Function-based Decorator

My basic plan is that we’ll have a decorator @report_creation that will log creations somehow, probably by just tallying a count or something. I really don’t know what’s going to happen when I wrap a class, so I start with this:

    def test_wrapping_class(self):
        def report_creation(*args, **kwargs):
            print(args, "*", kwargs)

        @report_creation
        class Reported:
            def __init__(self):
                pass

        reported = Reported()

I expect that something will print and then we’ll fail for some reason. Here’s what happens:

(<class 'test_decorators.TestDecorators.test_wrapping_class
.<locals>.Reported'>,) * {}
>       reported = Reported()
E       TypeError: 'NoneType' object is not callable

OK. Our function received the class Reported as its only parameter. The diagnostic just means that we didn’t return a function that could be called. Since it was the attempt to instantiate the class that failed, we can be pretty confident that our wrapper should return an instance of the wrapped class.

I’ll try this:

    def test_wrapping_class(self):
        def report_creation(klass):
            def wrapper_report():
                print("creating", klass)
                return klass
            return wrapper_report

        @report_creation
        class Reported:
            def __init__(self):
                pass

        reported = Reported()
        r2 = Reported()
        assert False

As I had hoped, this prints two “creating” lines:

creating <class 'test_decorators.TestDecorators
.test_wrapping_class.<locals>.Reported'>
creating <class 'test_decorators.TestDecorators
.test_wrapping_class.<locals>.Reported'>

We have learned that we should return the class to be created, in this case the original. I imagine that we could wrap that class in another class or function, but that’s beyond the scope of my thinking at this point. I wonder how we could change this wrapper to retain a count of instances created. It turns out that a function is an object and therefore a function can have attributes. Like this:

    def test_wrapping_class(self):
        def report_creation(klass):
            report_creation.count = 0
            def wrapper_report():
                report_creation.count += 1
                return klass
            return wrapper_report

        @report_creation
        class Reported:
            def __init__(self):
                pass

        reported = Reported()
        r2 = Reported()
        assert report_creation.count == 2

Now I think this one will be keeping each class count separately, but I’m not sure of that. Extend the test:

    def test_wrapping_class(self):
        def report_creation(klass):
            report_creation.count = 0
            def wrapper_report():
                report_creation.count += 1
                return klass
            return wrapper_report

        @report_creation
        class Reported:
            def __init__(self):
                pass

        @report_creation
        class Another:
            pass

        reported = Reported()
        r2 = Reported()
        another = Another()
        assert report_creation.count == 3

As you can see, I had to change the count assertion to 3, showing that the count applies to all the usages of the decorator function. I do not understand that.

It seems to me that the count is zeroed on every call to report creation. Ah. The test is working because I didn’t create any instances in between the two class definitions. Here’s a better test to show what happens:

        @report_creation
        class Reported:
            def __init__(self):
                pass

        reported = Reported()
        r2 = Reported()
        assert report_creation.count == 2

        @report_creation
        class Another:
            pass

        another = Another()
        assert report_creation.count == 1

That passes. Suppose we actually wanted the combined count. We could do this:

    def test_wrapping_class(self):
        def report_creation(klass):
            if not hasattr(report_creation, "count"):
                report_creation.count = 0
            def wrapper_report():
                report_creation.count += 1
                return klass
            return wrapper_report

        @report_creation
        class Reported:
            def __init__(self):
                pass

        reported = Reported()
        r2 = Reported()
        assert report_creation.count == 2

        @report_creation
        class Another:
            pass

        another = Another()
        assert report_creation.count == 3

So this is nice. We have a stateful decorator, implemented by stuffing an attribute into the function object, only if it isn’t already there. That’s getting a bit deep, but it should be clear that we could do something like that to implement Singleton pattern, if we wanted to. In that case, we’d want to keep the state variable separate for each class created, as in the earlier example.

Should we do Singleton, just to show that we know how, or just as another bit of learning? Sure, let’s do.

    def test_singleton(self):
        # lifted from RealPython
        def singleton(klass):
            def wrapper_singleton():
                if not wrapper_singleton.instance:
                    wrapper_singleton.instance = klass()
                return wrapper_singleton.instance
            wrapper_singleton.instance = None
            return wrapper_singleton

        @singleton
        class Something:
            pass

        s1 = Something()
        s2 = Something()
        assert s1 is s2

I state here in front of Hill and everyone that I peeked at an example of Singleton on RealPython before I did this, because I had been studying the decorator topic and I knew it was there. It contaminated my mind enough that I want to make clear that I did the above from memory of a solution, not by devising it. It shows up in the names if nowhere else.

Note
In both these examples of decorating classes, we would need to provide the standard *args*, **kwargs parameters to allow for all forms of class creation. I left those out to better show the form of the solution, but it would go like this:
    def test_singleton(self):
        # lifted from RealPython
        def singleton(klass):
            def wrapper_singleton(*args, **kwargs):
                if not wrapper_singleton.instance:
                    wrapper_singleton.instance = klass(*args, **kwargs)
                return wrapper_singleton.instance
            wrapper_singleton.instance = None
            return wrapper_singleton

        @singleton
        class Something:
            def __init__(self, a, b):
                self.a = a
                self.b = b


        s1 = Something(1, 2)
        s2 = Something(3, 4)
        assert s1 is s2
        assert s2.a == 1

Note that I added a check on the parameters passed, just for fun, showing that s2’s attempt to set the a and b didn’t take.

Reflection

There is still another thing that we could test, in principle, a parameterized decorator for decorating a class. We’ve done one that decorates a function, and determined that we need a second level of wrapper to manage that situation. I don’t plan to do that this morning, and perhaps never, because of this important fact:

Given a function foo and a class Bar, there is no essential difference between these two lines:

foo = foo()
bar = Bar()

A function is an object that you can call. And a class is also an object that you can call. When you call the function, any concrete parameters that you provide are passed to the function. When you call the class, any parameters are passed to the class’s __init__ method.

It follows that we would do a parameterized decorator for a class just exactly the same way we’d do one for a function.

Summary / Recipes

On the road to a recipe:

Whether we are decorating a function or a class, we’re really decorating a “callable”, an object that can be called. Therefore our decorator must return a callable. That callable could be the original function or class, a new function, or a new class: it just has to be callable.

No Decorator Parameters

In the case where the decorator itself has no parameters, we declare our decorator with a single parameter, the function to be called:

def my_decorator(func):

The decorator defines a new function which will wrap the input function, and returns that. The wrapping function must accept arbitrary arguments to be ready for any function. And the decorator returns that function:

def my_decorator(func):
    def my_wrapper(*args, **kwargs):
        ...
    return my_wrapper

The general rule will be always to use the *args,**kwargs form. Otherwise our decorator will be limited to whatever single form we use to call the function, when we get around to it.

Generally speaking, our decorator may do other things before or after calling the decorated function, but it will generally call it. (A caching decorator might only call the function on new inputs. Example probably left to the reader. Or you can find one on the Internet.) So our final decorator might look like this:

def my_decorator(func):
    def my_wrapper(*args, **kwargs):
        # do things before calling
        func(*args, **kwargs)
        # do things after calling
    return my_wrapper

Now, as we’ve seen in our @repeat example, the wrapper might call the function more than once and as we can imagine in the cached situation, it might call it less than once, but if and when it does call, it calls with the parameters as shown.

One more thing. As a nicety, it is recommended to use functools.wraps to help Python produce a more useful name for our wrapped function in the event of trouble. This is just part of the recipe. It’s optional but pythonic. Recommended:

def my_decorator(func):
    functools.wraps(func)
    def my_wrapper(*args, **kwargs):
        # do things before calling
        func(*args, **kwargs)
        # do things after calling
    return my_wrapper

With Decorator Parameters

It is possible to define a decorator that itself takes parameters. We could imagine a logging decorator that allowed us to include a header line for our log entries. And we have in today’s article an example of a decorator that takes a numeric parameter, used to say how often to call the wrapped function.

The major difference between a decorator with parameters and one without is that Everything Changes. Well, not quite. But we start with our decorator expecting its own parameters, not the function being wrapped:

def my_decorator(my_parameter):

The parameter(s) will be visible throughout in what follows. What follows next is that we provide an inner function that will be called to pass in the function to be wrapped:, and return that.

def my_decorator(my_parameter):
    def inner(func):
        ...
    return inner

The inner function will be called, passing in the actual user function to be wrapped, so we wrap it, using the wraps again to be nice, and return that:

def my_decorator(my_parameter):
    def inner(func):
        functools.wraps(func)
        def my_wrapper(*args, **kwargs):
            # do things before
            func(*args, **kwargs)
            # do things after
        return my_wrapper
    return inner

Here’s the essence, without decorator parameters and with them:

# without decorator parameters
def my_decorator(func):
    functools.wraps(func)
    def my_wrapper(*args, **kwargs):
        # do things before calling
        func(*args, **kwargs)
        # do things after calling
    return my_wrapper

# with decorator parameters
def my_decorator(my_parameter):
    def inner(func):
        functools.wraps(func)
        def my_wrapper(*args, **kwargs):
            # do things before calling
            func(*args, **kwargs)
            # do things after calling
        return my_wrapper
    return inner

Summary / So Far

Commit: end of article 263.

We have worked out how decorators work by writing tests. Early tests just accepted parameters and printed them, so that we could figure out what was being a decent understanding of the “recipe” for creating decorators using functions. Later tests extended those early ones, until we had tests running specific examples of specific kinds of decorators. Those tests can be reviewed to se how to build new decorators should we care to.

Prior to becoming test-addicted, I would have done similar things to work out how decorators worked, perhaps working from the command line, or running a small main program that printed things. I like this scheme much better, because it preserves my work as the growing set of tests. I also committed very frequently, though I didn’t report all the commits. Every time I made a little progress, I committed “save point” That way, if my next attempt interfered too badly with the canine, I had a safe place to roll back to.

I am quite pleased with how this approach allowed me to sneak up on understanding decorators. I feel fairly good about my understanding now, probably about an 85 out of 100. Not 100 as in “knows all things about decorators”. More like “knows what he wants to know for now”.

We’ll do the same to work out the recipe for creating them with classes. I expect that may feel more natural to me, but we’ll find out. I do anticipate that the two recipes will be quite similar. After all … a class is just another kind of callable thing.

Below please find the current version of all my tests. See you next time!


class TestDecorators:
    def test_hookup(self):
        assert 4 == 2+2

    def test_repeat_twice(self):
        def repeat(func):
            functools.wraps(func)
            def rpt(*args, **kwargs):
                func(*args, **kwargs)
                func(*args, **kwargs)
            return rpt
        count = 0

        @repeat
        def increment():
            nonlocal count
            count += 1
        assert count == 0
        increment()
        assert count == 2

        s = ""

        @repeat
        def add_a():
            nonlocal s
            s += "a"
        add_a()
        assert s == "aa"

    def test_repeat_with_parameter(self):
        def repeat(times):
            def rpt(func):
                @functools.wraps(func)
                def wrapped(*args, **kwargs):
                    for i in range(times):
                        func(*args, **kwargs)
                return wrapped
            return rpt

        count = 0

        @repeat(3)
        def increment():
            nonlocal count
            count += 1
        print("increment", increment)
        count = 0
        increment()
        assert count == 3

    def test_repeat_with_function_parameters(self):
        def repeat(times):
            def rpt(func):
                @functools.wraps(func)
                def wrapped(*args, **kwargs):
                    for i in range(times):
                        func(*args, **kwargs)
                return wrapped
            return rpt
        count = 0
        @repeat(3)
        def increment(by, multiplier=1):
            nonlocal count
            count += by*multiplier
        increment(2)
        assert count == 6
        count = 0
        increment(2, multiplier=2)
        assert count == 12

    def test_wrapping_class(self):
        # note that we do not actually wrap the class itself,
        # we just return the original. A more complex example
        # might do more.
        def report_creation(klass):
            if not hasattr(report_creation, "count"):
                report_creation.count = 0
            def wrapper_report():
                report_creation.count += 1
                return klass
            return wrapper_report

        @report_creation
        class Reported:
            def __init__(self):
                pass

        reported = Reported()
        r2 = Reported()
        assert report_creation.count == 2

        @report_creation
        class Another:
            pass

        another = Another()
        assert report_creation.count == 3

    def test_singleton(self):
        # lifted from RealPython
        def singleton(klass):
            def wrapper_singleton(*args, **kwargs):
                if not wrapper_singleton.instance:
                    wrapper_singleton.instance = klass(*args, **kwargs)
                return wrapper_singleton.instance
            wrapper_singleton.instance = None
            return wrapper_singleton

        @singleton
        class Something:
            def __init__(self, a, b):
                self.a = a
                self.b = b


        s1 = Something(1, 2)
        s2 = Something(3, 4)
        assert s1 is s2
        assert s2.a == 1