Ivan Sanych
Ivan Sanych

Reputation: 27

How do python contexts relate to variable scoping?

I read examples from ""Bayesian Methods for Hackers" and I noticed a sort of "magic" obtained with python context:

In PyMC3, we typically handle all the variables we want in our model within the context of the Model object. This is an extra layer of convenience compared to PyMC. Any variables created within a given Model's context will be automatically assigned to that model. If you try to define a variable outside of the context of a model, yo u will get an error.

The code in question is this

import pymc3 as pm

with pm.Model() as model:
    parameter = pm.Exponential("poisson_param", 1.0)
    data_generator = pm.Poisson("data_generator", parameter)

I tried to replicate this behavior with a custom toy class that defines __enter__ and __exit__ methods but couldn't create a class which works such that "Any variables created within a given Model's context will be automatically assigned to that model"

Upvotes: 1

Views: 77

Answers (1)

AKX
AKX

Reputation: 169417

Sure.

The objects are basically "smart" in that they know somehow what the current context (model) is and register themselves within it.

Here's an example (that also supports entering nested contexts, though they currently don't have any sort of parent/child hierarchy among themselves).

class Context:
    # Stack of current contexts entered; stored as a class attribute
    stack = []

    def __init__(self, name):
        self.name = name
        self.vars = {}

    def __enter__(self):
        # Put self on the stack.
        Context.stack.append(self)
        return self

    def __exit__(self, et, ev, tb):
        # Ensure we're popping ourselves off the stack.
        assert Context.stack[-1] is self
        # Remove self.
        Context.stack.pop(-1)

    def register(self, var):
        # Register a `Var` in self...
        self.vars[var.name] = var
        # ... and set its Context.
        var.context = self

    @staticmethod
    def get_current():
        # Get the topmost Context if any.
        return Context.stack[-1] if Context.stack else None

    def __repr__(self):
        return "<Context %r with %d vars>" % (self.name, len(self.vars))


class Var:
    def __init__(self, name):
        self.context = None
        self.name = name

        # Register ourselves in the current context, if any.
        ctx = Context.get_current()
        if ctx:
            ctx.register(self)

    def __repr__(self):
        return "<%s (in %s)>" % (self.name, self.context)


with Context("spam") as c:
    v = Var("foo")
    x = Var("bar")

print(c.vars)

Upvotes: 1

Related Questions