How can we "associate" a Python context manager to the variables appearing in its block?

Question

As I understand it, context managers are used in Python for defining initializing and finalizing pieces of code (__enter__ and __exit__) for an object.

However, in the tutorial for PyMC3 they show the following context manager example:

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)

and mention that this has the purpose of associating the variables alpha, beta, sigma, mu and Y_obs to the model basic_model.

I would like to understand how such a mechanism works. In the explanations of context managers I have found, I did not see anything suggesting how variables or objects defined within the context's block get somehow "associated" to the context manager. It would seem that the library (PyMC3) somehow has access to the "current" context manager so it can associate each newly created statement to it behind the scenes. But how can the library get access to the context manager?

Answer 1

PyMC3 does this by maintaining a thread local variable as a class variable inside the Context class. Models inherit from Context.

Each time you call with on a model, the current model gets pushed onto the thread-specific context stack. The top of the stack thus always refers to the innermost (most recent) model used as a context manager.

Contexts (and thus Models) have a .get_context() class method to obtain the top of the context stack.

Distributions call Model.get_context() when they are created to associate themselves with the innermost model.

So in short:

1. with model pushes model onto the context stack. This means that inside of the with block, type(model).contexts or Model.contexts, or Context.contexts now contain model as its last (top-most) element.
2. Distribution.__init__() calls Model.get_context() (note capital M), which returns the top of the context stack. In our case this is model. The context stack is thread-local (there is one per thread), but it is not instance-specific. If there is only a single thread, there also is only a single context stack, regardless of the number of models.
3. When exiting the context manager. model gets popped from the context stack.