Deep Python #1: With magic

What are the "with" keyword and context managers, why they are not limited to system resources, and great code examples using them.

Introduction: with, context managers

“With” is one of the most pythonic statements out there. It is a keyword and a concept you do not find easily in other mainstream languages. It is used with context managers, which at their essence are just objects implementing the __enter__() and __exit__() methods, which in turn are automatically called by with.

The intended goal of with and context managers is to manage resources while writing as little scaffolding as possible. The most typical example is

with open(filename, r_or_w) as f:
    # do stuff with f

# more code that does not require f

In this example, open is a context manager whose __enter__() method returns the open file handler, here called f, while its __exit__() method closes it. When using a with statement, __enter__() is automatically called after creating the context manager object, and __exit__() is called when the context manager goes out of scope. It is called regardless of the reason why it goes out of scope, by normal code execution or by raising an exception. This way, we know that we are properly cleaning up after ourselves. It saves us from writing complex cleaning up logic to take into account all cases. It just works.

Some resources that should normally be closed to avoid leaks are: files (like in the example above), sockets, database connections, processes, threads. All of them borrow system resources that must be returned to prevent system degradation.

Were it just about this, with statements would only be a shorter form of a try … finally code block to correctly manage resources, but since we can write any code we want inside the __enter__() and __exit__() methods, they can empower us to write elegant and powerful code for many different tasks.

Tasting the magic

Before starting with real code, let us see a short toy example of how a context manager can be used to change code behavior in a limited scope:

import random


class Fixed:
    def __enter__(self):
        self.r = random.random
        random.random = lambda: 5  # when called random() return 5
        return "NOT RANDOM"   # just something to return in case 'as' is used

    def __exit__(self, exc_type, exc_val, exc_tb):
        random.random = self.r
        # the arguments allow us to do something with raised exceptions, but we don't care for this example


if __name__ == '__main__':
    print([random.random() for _ in range(10)])
    with Fixed():
        print([random.random() for _ in range(10)])

    print([random.random() for _ in range(10)])

The above code, which you can run as it is, uses the Fixed context manager to change locally the behavior of random.random(). It replaces random.random() with a function that takes no arguments and always returns the number 5 when the context is entered, and the original function is restored at exit.

Example output:

[0.8443463505724088, 0.5375917438523139, 0.12024430301030031, 0.10678952662762009, 0.966354234275387, 0.5620289999783581, 0.1890188347213756, 0.33537191743312356, 0.6151066813028078, 0.6038920031969371]

[5, 5, 5, 5, 5, 5, 5, 5, 5, 5]

[0.3287682665937276, 0.8867162956099921, 0.789230280474005, 0.19740073039418227, 0.5594914687478751, 0.1468699258846229, 0.9374862487775815, 0.7193982499414655, 0.7245365001109647, 0.7657307069285624]

and you can easily see that using the with keyword makes it very clear and explicit, with no complex code for the user.

As for all magic tricks, the magic actually vanishes when the trick is explained. Yet, being able to create this kind of abstraction is a useful skill for Python developers and it is indeed widely applied in many real code bases.

It is now time for some real code!

PyTorch no_grad

Pytorch is a deep learning framework that, among many other things, allows us to define a computation graph that can be executed forward (as we defined it) and backwards by following a second computation graph based on the gradients that is built automatically. This is particularly useful for training models with little scaffolding code.

torch.no_grad() is a context manager to deactivate gradient computation for some tensor variables. It is used during inference to avoid useless computation overhead as gradients are (usually) only useful during training.

Example from the docs:

>>> x = torch.tensor([1.], requires_grad=True)
>>> with torch.no_grad():
...     y = x * 2
>>> y.requires_grad
False

so no gradients will be computed for y, while the normal behavior would be to compute it. This little magic, again very readable, is done with very few lines of code. From the source:

class no_grad(_NoParamDecoratorContextManager):
    """ docs ... """
    def __init__(self) -> None:
        if not torch._jit_internal.is_scripting():
            super().__init__()
        self.prev = False

    def __enter__(self) -> None:
        self.prev = torch.is_grad_enabled()
        torch.set_grad_enabled(False)

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        torch.set_grad_enabled(self.prev)

Pytorch has functions to toggle gradient computation and get its current value. This context manager just replaces that value temporarily like we did in our toy example above, and then goes back to the previous state in __exit__().

PuePy tag nesting

PuePy is an exciting new framework to write frontend code entirely in python, by leveraging pyscript and webassembly. In the project homepage we see this code sample to generate html from python:

from puepy import Page, t

class CounterPage(Page):
    def initial(self):
        return {"current_value": 0}

    def populate(self):
        with t.div(classes="inc-box"):
            t.button("-", on_click=self.on_decr)
            t.span(str(self.state["current_value"]))
            t.button("+", on_click=self.on_incr)

    def on_decr(self, event):
        self.state["current_value"] -= 1

    def on_incr(self, event):
        self.state["current_value"] +=  1 				

it basically creates a page with two buttons and a number, and the two buttons respectively decrease and increase it.

In this code sample a context manager is used to nest the two buttons and a span tags into a div tag. In puepy, t is a builder for html tags. Each tag is an instance of the Tag class and this is what we are going to look at:

class Tag:
    """
    A class that allows you to push and pop a context
    """

    stack = []
    population_stack = []
    origin_stack = [[]]
    component_stack = []
    default_classes = []
    default_attrs = {}
    default_role = None

    document = document

    def __init__( 
       .
       .
       .

    def __enter__(self):
        self.stack.append(self)
        self.origin_stack[0].append(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stack.pop()
        self.origin_stack[0].pop()
        return False

This time another trick is used: a class attribute to share memory between all instances of the Tag class. The Tag class also has a parent instance attribute that is initialized to the last element in the stack: the last used with determines the parent. Finally, the parent information is used to generate the html correctly.

This looks really, but may suffer from a minor bug: what if there are multiple threads accessing the same class attribute at the same time?

PyMC models

PyMC is a library for probabilistic programming that lets you define Bayesian models for your data and then run inference on them. Two important concepts in the library are Model and Distribution. A Distribution is a parametrized probability distribution (e.g. Gaussian, Bernoulli, Uniform, etc…), while a Model stores distributions as prior probabilities and likelihood for your data, and enables us to do some reasoning on them.

An example from their tutorial to make it more concrete:

basic_model = pm.Model()

with basic_model:

    # Priors for unknown model parameters
    alpha = pm.Normal('alpha', mu=0, sd=10)
    beta = pm.Normal('beta', mu=0, sd=10, shape=2)
    sigma = pm.HalfNormal('sigma', sd=1)

    # Expected value of outcome
    mu = alpha + beta[0]*X1 + beta[1]*X2

    # Likelihood (sampling distribution) of observations
    Y_obs = pm.Normal('Y_obs', mu=mu, sd=sigma, observed=Y)

everything inside the context is a Distribution. The model (basic_model) is never used explicitly, but a Distribution outside a Model raises an exception when it is initialized. A distribution must be defined inside a model context. So, how is this relationship written down in code?

Please, bear with me because this final example is the most complex so far.

The Model class is defined as:

class Model(WithMemoization, metaclass=ContextMeta):
.
.
.

and no __enter__() and __exit__() methods are present in it. To find them we have to check its metaclass ContextMeta:

class ContextMeta(type):
    """Functionality for objects that put themselves in a context using
    the `with` statement.
    """

    def __new__(cls, name, bases, dct, **kwargs):
        """Add __enter__ and __exit__ methods to the class."""

        def __enter__(self):
            self.__class__.context_class.get_contexts().append(self)
            return self

        def __exit__(self, typ, value, traceback):
            self.__class__.context_class.get_contexts().pop()

        dct[__enter__.__name__] = __enter__
        dct[__exit__.__name__] = __exit__

        # We strip off keyword args, per the warning from
        # StackExchange:
        # DO NOT send "**kwargs" to "type.__new__".  It won't catch them and
        # you'll get a "TypeError: type() takes 1 or 3 arguments" exception.
        return super().__new__(cls, name, bases, dct)

ContextMeta adds the two methods when its classes are created. But what do these methods actually do?

__enter__() adds itself to the class attribute context_class’s stack, returned by get_contexts().

def get_contexts(cls) -> list[T]:
        """ docs """
        context_class = cls.context_class
        assert isinstance(
            context_class, type
        ), f"Name of context class, {context_class} was not resolvable to a class"
        if not hasattr(context_class, "contexts"):
            context_class.contexts = threading.local()

        contexts = context_class.contexts

        if not hasattr(contexts, "stack"):
            contexts.stack = []
        return contexts.stack

the important parts are the end (return contexts.stack) and context_class.contexts = threading.local() which makes context_class.contexts a storage shared in the local thread, to make it thread-safe, unlike the previous code.

Finally, when a Distribution is constructed, it first uses Model.get_context() to get the latest model present in the stack of the Model class. Please note the capital M, we can access class attributes without having direct access to an instance.

try:
   from pymc.model import Model

   model = Model.get_context()
   except TypeError:
       raise TypeError(
           "No model on context stack, which is needed to "
           "instantiate distributions. Add variable inside "
           "a 'with model:' block, or use the '.dist' syntax "
           "for a standalone distribution."
       )

and then register itself to the model

        rv_out = cls.dist(*args, **kwargs)

        rv_out = model.register_rv(
            rv_out,
            name,
            observed=observed,
            total_size=total_size,
            dims=dims,
            transform=transform,
            default_transform=default_transform,
            initval=initval,
        )

        # add in pretty-printing support
        rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
        rv_out._repr_latex_ = types.MethodType(
            functools.partial(str_for_dist, formatting="latex"), rv_out
        )

        rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
        rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
        rv_out.random = _make_nice_attr_error("rv.random()", "pm.draw(rv)")
        return rv_out

but this complex mechanism is all hidden behind the simplicity of using
with pm.Model().

Conclusion

Context managers and the with statement are pythonic tools that is worth mastering to write simple and elegant code, which reads easily because it clearly shows its intent.

Though they were initially designed for resource management, we can use them every time we want to make changes limited to a scope.

In this post we have studied some methods used in real codebases, with different complexities, to make a creative use of the with statement. Hopefully, it can help you to write better Python code that makes you and your colleagues happy!