Python decorators and context managers

cache.py

def cache(obj):
    saved = obj.saved = {}
    @functools.wraps(obj)
    def memoizer(*args, **kwargs):
        key = str(args) + str(kwargs)
        if key not in saved:
            saved[key] = obj(*args, **kwargs)
        return saved[key]
    return memoizer

# now our nice clean web_lookup()

@cache
def web_lookup(url):
    page = urlopen(url)
    try:
        return page.read()
    finally:
        page.close() 

closing.py

from contextlib import closing

def web_lookup(url):
    with closing(urlopen(url)) as page:
        return page.read()

decorator_example.py

@cache
def web_lookup(url):
    page = urlopen(url)
    try:        
        return page.read()
    finally:
        page.close()

DecoratorsAndContentManagers.html

<p>
    Two weeks ago I wrote a post called <a href="https://codefisher.org/catch/blog/2015/01/27/python-tips-tricks-and-idioms/">Python: tips, tricks and idioms</a> where I went though a whole lot of features of python.  Now however I want to narrow down on just a few, and look at them in more depth.  The first is decorators, which I did not cover at all, and the second is context managers which I only gave one example of.
</p>

<p>
    There is a reason that I put them together; they both have the same goal.  They can both help separate what your trying to do (the "business logic") from some extra code added for clean up or performance etc. (the "administrative logic").  So basically it helps package away in a reusable way code that really we don't care too much about.
</p>

<h2>Decorators</h2>

<p>
    Decorators are really easy to use, and even if you have never seen them before. You could probably guess what is going on, even if not how or why.  Take this for example:
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=decorator_example.py"></script>

<p>
    Over looking for now exactly what library urlopen comes from... We can guess that the results of the <code>web_lookup</code> are going to be cached, so that they are not fetched ever time we ask for the same url.  So simple, just know <code>@some_decorator</code> before our function and we can use any decorator.  But how do we write one?  
</p>

<p>
    First we need to understand what the decorator is doing.  <code>@cache</code> is actually just syntactic sugar for the following.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=function_wrapped.py"></script>

<p>
    So this is important.  What cache is, is a function that takes another function as an argument, and returns a new function, that can be used just like the function could be before, but presumably adding in some extra logic.  So for our first decorator let us start with something simple, a decorator that squares the result of the function it wraps. 
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=square.py"></script>

<p>
    So in this little example every time we call <code>plus()</code> the number will have 1 added it it, but because of our decorator, the result will also be squared.
</p>

<p>
    But there is a problem with this, <code>plus()</code> is no longer the <code>plus()</code> function that we defined, but another function wrapping it.  So things like the doc string have gone missing.  Things like <code>help(plus)</code> will no longer work.  But in the functools library there is a decorator to fix that, <code>functools.wraps()</code>.  Always use <code>functools.wraps()</code> when writing a decorator.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=square_wraps.py"></script>

<p>
    But did you notice?  <code>wraps()</code> is a decorator that takes arguments.  How do we write something like that? It gets a little more involved, but let us start with the code.  It will be a function the now raises the number to some power.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=power.py"></script>

<p>
    So yes three functions, one inside the other.  The result now, is that <code>power()</code> is not longer really the decorator, but a function that returns the decorator.  See the issue is one of scoping, which we why we have to put the functions inside each other.  When <code>_pow()</code> is called, the value of <code>pow</code>, comes from the scope of the <code>power()</code> function that contains it.
</p>

<p>
    So now we know how to write highly reusable function decorators, or do we?  There is a problem still, and that is our internal function <code>_square()</code> or <code>_pow()</code> only takes one argument, so that any function it wraps, can only take one argument as well.  What we want is to be able to have any number of arguments.  So that is where the star operator comes in.
</p>

<h2>Star operator</h2>

<p>
    The * (star) operator can be used in a function definition so that it can take an arbitrary number of arguments, all of which are collapsed into a single tuple.  An example might help.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=star.py"></script>

<p>
    The * operator can also be used for the reverse case, when we have a itterator, but want to pass that as the arguments to a function.  This gets called argument unpacking.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=unpack.py"></script>

<p>
    The same basic idea, can also be used for keyword arguments, but for this we use the ** (double star) operator.  But instead of getting a list of the arguments, we get a dictionary.  We can also use both together.  So some examples.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=kwargs.py"></script>

<h2>Better Decorators</h2>

<p>
    So now we can go back and change our decorator to be truly generic.  Lets do it with the simplest one, we wrote, <code>@square</code>.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=square_star.py"></script>

<p>
    Now no matter what arguments the function takes, we will happily just pass them though to the function that we are wrapping.
</p>

<p>
    So let us go back to our <code>web_lookup</code> function, and write first it first with caching, and then write the decorator to see the difference.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=web_lookup.py"></script>

<p>
    That is how it might look, and our problem here is that the caching code is mixed up in with what <code>web_lookup()</code> is supposed to do.  Makes it harder to maintain it, reuse it, and also harder if we want to update the way we cache it if we have done something like this all over our code.  So our very generic decorator might look like this.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=cache.py"></script>


<p>
    So that can wrap any function with any number of arguments just by putting <code>@cache</code> before it.  But I did not write that function myself, I just lifted it right off the <a href="https://wiki.python.org/moin/PythonDecoratorLibrary">Python Decorator Library</a> that has many many examples of decorators you can use.  
</p>

<h2>Context Managers</h2>

<p>
    In the previous post I did a single example of using a context manager, opening a file.  It looked like this:
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=open_with.py"></script>

<p>
    Admittedly the context manager is only a little shorter, and the file will be garbage collected (at least in CPython) but there are other cases where it might be a lot bigger problem if you get it wrong.  For example the threading library can also use a context manager.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=lock.py"></script>

<p>
    Which is nice and simple, and you can see from the indent what the critical section is, and also be sure the lock is released.
</p>

<p>
    If you dealing with file like objects that don't have a content manager, the <a href="https://docs.python.org/3.4/library/contextlib.html">contextlib</a> has the closing context manager.  So let us go back an improve our <code>web_lookup()</code> function.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=closing.py"></script>

<p>
    We can also write our own context managers.  All that is needed to to used the <code>@contextmanager</code> decorator, and have a function with a <code>yeild</code> in it.  The <code>yeild</code> marks the point at which the context manager stops while the code with in the <code>with</code> statement runs.  So the following can be used to time how long it takes to do something.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=timeit.py"></script>

<p>
    The try finally in this case is kind of optional, but without it the time would not be printed if there was an exception raised inside the with statment.
</p>

<p>
    We can also do more complicated context managers.  The following is something like what was added in Python 3.4.  It will take what every is printed to sysout, and put it in a file (or file like object).  So for example if we had all our <code>timeit()</code> context managers in your code, and wanted to start putting the results into a log file.  Here the <code>yeild</code> is followed by a value, which is why we can then use the <code>with ... as</code> syntax.
</p>

<script src="https://gist.github.com/codefisher/af24fe8b1da830ad3f43.js?file=redirect_stdout.py"></script>


<p>
    The last with statement also shows off the use of compound <code>with</code> statements.  It is really just the same as putting one <code>with</code> inside another.
</p>

<p>
    Finally it is worth mentioning at least in passing, that any <code>class</code> can be turned into a context manager by adding the <code>__enter__()</code> and <code>__exit__()</code> methods.  The code in either will do more or less what the code either side of the yield statement would do.
</p>

<p>
    And that is all for this round, hope you learned something new and interesting.  Don't forget to follow me on twitter if you want more python tips, such as when next time I write about sets and dictionaries.  In the mean time if you looking for more there is a great book <a href="http://www.amazon.com/gp/product/1449340377/ref=as_li_tl?ie=UTF8&camp=1789&creative=9325&creativeASIN=1449340377&linkCode=as2&tag=codefisherorg-20&linkId=BRTKQK5IF4JXR6OL">Python Cookbook, Third edition</a><img src="http://ir-na.amazon-adsystem.com/e/ir?t=codefisherorg-20&l=as2&o=1&a=1449340377" width="1" height="1" alt="" style="border:none !important; margin:0px !important;" /> by O'Reilly Media.  I have been reading parts of it, and might include a few things I learned from it in my next post.  Or if you want something simpler try <a href="http://www.amazon.com/gp/product/1449355730/ref=as_li_tl?ie=UTF8&camp=1789&creative=9325&creativeASIN=1449355730&linkCode=as2&tag=codefisherorg-20&linkId=FMW3BECBN3KH3N6U">Learning Python, 5th Edition</a><img src="http://ir-na.amazon-adsystem.com/e/ir?t=codefisherorg-20&l=as2&o=1&a=1449355730" width="1" height="1" alt="" style="border:none !important; margin:0px !important;" />.
</p>

function_wrapped.py

web_lookup = cache(web_lookup)

kwargs.py

def print_args(*args, **kwargs):
    print(args)
    print(kwargs)

print_args("Hello", "world", count=2, letters=10)

# output:
# ('Hello', 'world')
# {'count': 2, 'letters': 10}

# or calling the function with argument unpacking.

words = ("Hello", "world")
arguments = {'count': 2, 'letters': 10}

print_args(*words, **arguments)

lock.py

lock = threading.Lock()
with lock:
    print('Critical section')

open_with.py

with open('/etc/passwd', 'r') as f:
    print(f.read())

# which is equivalent to the longer

f = open('/etc/passwd', 'r')
try:
    print(f.read())
finally:
    f.close()

power.py

def power(pow):
    def _power(func):
        @wraps(func)
        def _pow(num):
            return func(num) ** pow
        return _pow 
    return _power

@power(2) 
def plus(num):
    """Adds 1 to a number"""
    return num + 1

redirect_stdout.py

from contextlib import contextmanager
import io, sys

@contextmanager
def redirect_stdout(fileobj=None):
    if fileobj is None:
        fileobj = io.StringIO() # in python 2 use BytesIO
    oldstdout = sys.stdout
    sys.stdout = fileobj
    try:
        yield fileobj
    finally:
        sys.stdout = oldstdout

with redirect_stdout() as f:
    help(pow)
help_text = f.getvalue()

with open('some_log_file', 'w') as f:
    with redirect_stdout(f):
        help(pow)

# above can be also written as
with open('some_log_file', 'w') as f, redirect_stdout(f):
    help(pow)

square.py

def square(func):
    def _square(num):
        return func(num) ** 2
    return _square

# which we can then use like this

@square 
def plus(num):
    """Adds 1 to a number"""
    return num + 1

square_star.py

def square(func):
    @wraps(func)
    def __square(*args, **kwargs):
        return func(*args, **kwargs) ** 2
    return __square

square_wraps.py

from functools import wraps

def square(func):
    @wraps(func)
    def _square(num):
        return func(num) ** 2
    return _square

star.py

def join_words(*args):
    """Joins all the words into a single string"""
    return " ".join(args)

print(join_words("Hello", "world"))

timeit.py

from contextlib import contextmanager
import time

@contextmanager
def timeit():
    start = time.time()
    try:
        yield
    finally:
        print("It took", time.time() - start, "seconds")

# this might take a few seconds
with timeit():
    list(range(1000000))

unpack.py

words = ("Hello", "world")
print(join_words(*words))

web_lookup.py

saved = {}

def web_lookup(url):
    if url in saved:
        return saved[url]
    page = urlopen(url)
    try:        
        data = page.read()
    finally:
        page.close()
    saved[url] = data
    return data