Adding Decorator Arguments

The timer decorator of the prior section works, but it would be nice if it was more configurable—providing an output label and turning trace messages on and off, for instance, might be useful in a general-purpose tool like this. Decorator arguments come in handy here: when they’re coded properly, we can use them to specify configuration options that can vary for each decorated function. A label, for instance, might be added as follows:

def timer(label=''):
    def decorator(func):
        def onCall(*args):          # args passed to function
            ...                     # func retained in enclosing scope
            print(label, ...        # label retained in enclosing scope
        return onCall
    return decorator                # Returns that actual decorator

@timer('==>')                       # Like listcomp = timer('==>')(listcomp)
def listcomp(N): ...                # listcomp is rebound to decorator

listcomp(...)                       # Really calls decorator

广告:个人专属 VPN,独立 IP,无限流量,多机房切换,还可以屏蔽广告和恶意软件,每月最低仅 5 美元

This code adds an enclosing scope to retain a decorator argument for use on a later actual call. When the listcomp function is defined, it really invokes decorator (the result of timer, run before decoration actually occurs), with the label value available in its enclosing scope. That is, timer returns the decorator, which remembers both the decorator argument and the original function and returns a callable which invokes the original function on later calls.

We can put this structure to use in our timer to allow a label and a trace control flag to be passed in at decoration time. Here’s an example that does just that, coded in a module file named mytools.py so it can be imported as a general tool:

import time

def timer(label='', trace=True):                  # On decorator args: retain args
    class Timer:
        def __init__(self, func):                 # On @: retain decorated func
            self.func    = func
            self.alltime = 0
        def __call__(self, *args, **kargs):       # On calls: call original
            start   = time.clock()
            result  = self.func(*args, **kargs)
            elapsed = time.clock() - start
            self.alltime += elapsed
            if trace:
                format = '%s %s: %.5f, %.5f'
                values = (label, self.func.__name__, elapsed, self.alltime)
                print(format % values)
            return result
    return Timer

Mostly all we’ve done here is embed the original Timer class in an enclosing function, in order to create a scope that retains the decorator arguments. The outer timer function is called before decoration occurs, and it simply returns the Timer class to serve as the actual decorator. On decoration, an instance of Timer is made that remembers the decorated function itself, but also has access to the decorator arguments in the enclosing function scope.

This time, rather than embedding self-test code in this file, we’ll run the decorator in a different file. Here’s a client of our timer decorator, the module file testseqs.py, applying it to sequence iteration alternatives again:

from mytools import timer

@timer(label='[CCC]==>')
def listcomp(N):                           # Like listcomp = timer(...)(listcomp)
    return [x * 2 for x in range(N)]       # listcomp(...) triggers Timer.__call__

@timer(trace=True, label='[MMM]==>')
def mapcall(N):
    return map((lambda x: x * 2), range(N))

for func in (listcomp, mapcall):
    print('')
    result = func(5)        # Time for this call, all calls, return value
    func(50000)
    func(500000)
    func(1000000)
    print(result)
    print('allTime = %s' % func.alltime)   # Total time for all calls

print('map/comp = %s' % round(mapcall.alltime / listcomp.alltime, 3))

Again, if you wish to run this fairly in 3.0, wrap the map function in a list call. When run as-is in 2.6, this file prints the following output—each decorated function now has a label of its own, defined by decorator arguments:

[CCC]==> listcomp: 0.00003, 0.00003
[CCC]==> listcomp: 0.00640, 0.00643
[CCC]==> listcomp: 0.08687, 0.09330
[CCC]==> listcomp: 0.17911, 0.27241
[0, 2, 4, 6, 8]
allTime = 0.272407666337

[MMM]==> mapcall: 0.00004, 0.00004
[MMM]==> mapcall: 0.01340, 0.01343
[MMM]==> mapcall: 0.13907, 0.15250
[MMM]==> mapcall: 0.27907, 0.43157
[0, 2, 4, 6, 8]
allTime = 0.431572169089
map/comp = 1.584

As usual, we can also test this interactively to see how the configuration arguments come into play:

>>> from mytools import timer
>>> @timer(trace=False)                    # No tracing, collect total time
... def listcomp(N):
...     return [x * 2 for x in range(N)]
...
>>> x = listcomp(5000)
>>> x = listcomp(5000)
>>> x = listcomp(5000)
>>> listcomp
<mytools.Timer instance at 0x025C77B0>
>>> listcomp.alltime
0.0051938863738243413

>>> @timer(trace=True, label='\t=>')       # Turn on tracing
... def listcomp(N):
...     return [x * 2 for x in range(N)]
...
>>> x = listcomp(5000)
        => listcomp: 0.00155, 0.00155
>>> x = listcomp(5000)
        => listcomp: 0.00156, 0.00311
>>> x = listcomp(5000)
        => listcomp: 0.00174, 0.00486
>>> listcomp.alltime
0.0048562736325408196

This timing function decorator can be used for any function, both in modules and interactively. In other words, it automatically qualifies as a general-purpose tool for timing code in our scripts. Watch for another example of decorator arguments in the section Implementing Private Attributes, and again in A Basic Range-Testing Decorator for Positional Arguments.


Note

Timing methods: This section’s timer decorator works on any function, but a minor rewrite is required to be able to apply it to class methods too. In short, as our earlier section Class Blunders I: Decorating Class Methods illustrated, it must avoid using a nested class. Because this mutation will be a subject of one of our end-of-chapter quiz questions, though, I’ll avoid giving away the answer completely here.