Metaclass-based solution.

Alternative approach to solving the context stack issues.
The previous version was a dead end, because there was no effective way to check the appropriateness of the classes on a single context stack.
So in this version, I split the context stacks into two, one for the pm.Models, and one for pm.distributions.distributions._DrawValuesContext.
This works better, but involved replacing the Context *class* with the ContextMeta parameterized metaclass.

Author: rpgoldman

pymc3/distributions/distribution.py

@@ -7,7 +7,7 @@
 from ..memoize import memoize
 from ..model import (
     Model, get_named_nodes_and_relations, FreeRV,
-    ObservedRV, MultiObservedRV, Context, InitContextMeta
+    ObservedRV, MultiObservedRV, ContextMeta
 )
 from ..vartypes import string_types, theano_constant
 from .shape_utils import (
@@ -449,7 +449,7 @@ def random(self, point=None, size=None, **kwargs):
                             "Define a custom random method and pass it as kwarg random")
 
 
-class _DrawValuesContext(Context, metaclass=InitContextMeta):
+class _DrawValuesContext(metaclass=ContextMeta, context_class='_DrawValuesContext'):
     """ A context manager class used while drawing values with draw_values
     """
 
@@ -476,7 +476,7 @@ def parent(self):
         return self._parent
 
 
-class _DrawValuesContextBlocker(_DrawValuesContext, metaclass=InitContextMeta):
+class _DrawValuesContextBlocker(_DrawValuesContext):
     """
     Context manager that starts a new drawn variables context disregarding all
     parent contexts. This can be used inside a random method to ensure that

pymc3/model.py

@@ -3,7 +3,8 @@
 import itertools
 import threading
 import warnings
-from typing import Optional, Tuple, TypeVar, Type, List
+from typing import Optional, Tuple, TypeVar, Type, List, Union
+from sys import modules
 
 import numpy as np
 from pandas import Series
@@ -162,59 +163,103 @@ def _get_named_nodes_and_relations(graph, parent, leaf_nodes,
             node_children.update(temp_tree)
     return leaf_nodes, node_parents, node_children
 
-T = TypeVar('T', bound='Context')
+T = TypeVar('T', bound='ContextMeta')
 
-class Context:
+
+class ContextMeta(type):
     """Functionality for objects that put themselves in a context using
     the `with` statement.
     """
-    contexts = threading.local()
-
-    def __enter__(self):
-        type(self).get_contexts().append(self)
-        # self._theano_config is set in Model.__new__
-        if hasattr(self, '_theano_config'):
-            self._old_theano_config = set_theano_conf(self._theano_config)
-        return self
-
-    def __exit__(self, typ, value, traceback):
-        type(self).get_contexts().pop()
-        # self._theano_config is set in Model.__new__
-        if hasattr(self, '_old_theano_config'):
-            set_theano_conf(self._old_theano_config)
-
-    @classmethod
-    def get_contexts(cls: Type[T]) -> List[T]:
-        # no race-condition here, cls.contexts is a thread-local object
-        # be sure not to override contexts in a subclass however!
-        if not hasattr(cls.contexts, 'stack'):
-            cls.contexts.stack = []
-        return cls.contexts.stack
-
-    @classmethod
-    def get_context(cls: Type[T], error_if_none=True) -> Optional[T]:
+    _context_class = None # type: Union[Type, str]
+
+    def __new__(cls, name, bases, dct,  **kargs):
+        # DO NOT send "**kargs" to "type.__new__".  It won't catch them and
+        # you'll get a "TypeError: type() takes 1 or 3 arguments" exception.   
+        # dct['get_context'] = classmethod(_get_context)
+        # dct['get_contexts'] = classmethod(_get_contexts)
+        return super().__new__(cls, name, bases, dct)
+
+    def __init__(cls, name, bases, nmspc, context_class: Optional[Type]=None, **kwargs):
+        if context_class is not None:
+            cls._context_class = context_class
+        super().__init__(name, bases, nmspc)
+        cls.contexts = threading.local()
+        def __enter__(self):
+            self.__class__.context_class.get_contexts().append(self)
+            # self._theano_config is set in Model.__new__
+            if hasattr(self, '_theano_config'):
+                self._old_theano_config = set_theano_conf(self._theano_config)
+            return self
+
+        def __exit__(self, typ, value, traceback):
+            self.__class__.context_class.get_contexts().pop()
+            # self._theano_config is set in Model.__new__
+            if hasattr(self, '_old_theano_config'):
+                set_theano_conf(self._old_theano_config)
+
+        cls.__enter__ = __enter__
+        cls.__exit__ = __exit__
+
+    def get_context(cls, error_if_none=True) -> Optional[T]:
         """Return the most recently pushed context object of type ``cls``
         on the stack, or ``None``."""
         idx = -1
         while True:
             try:
-                candidate = cls.get_contexts()[idx]
+                candidate = cls.get_contexts()[idx] # type: Optional[T]
             except IndexError as e:
+                # Calling code expects to get a TypeError if the entity
+                # is unfound, and there's too much to fix.
                 if error_if_none:
-                    raise e
+                    raise TypeError("No %s on context stack"%str(cls))
                 return None
-            if isinstance(candidate, cls):
-                return candidate
+            return candidate
             idx = idx - 1
 
+    def get_contexts(cls) -> List[T]:
+        # no race-condition here, cls.contexts is a thread-local object
+        # be sure not to override contexts in a subclass however!
+        if not hasattr(cls.context_class, 'stack'):
+            cls.context_class.stack = []
+        return cls.context_class.stack
+
+    @property
+    def context_class(cls) -> Type:
+        def resolve_type(c: Union[Type, str]) -> Type:
+            if isinstance(c, str):
+                c = getattr(modules[cls.__module__], c)
+            if isinstance(c, type):
+                return c
+            raise ValueError("Cannot resolve context class %s"%c)
+        assert cls is not None
+        if isinstance(cls._context_class, str):
+            cls._context_class = resolve_type(cls._context_class)
+        if not isinstance(cls._context_class, (str, type)):
+            raise ValueError("Context class for %s, %s, is not of the right type"%\
+                             (cls.__name__, cls._context_class))
+        return cls._context_class
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.context_class = super().context_class
+
+    # Initialize object in its own context...
+    def __call__(cls, *args, **kwargs):
+        instance = cls.__new__(cls, *args, **kwargs)
+        with instance:  # appends context
+            instance.__init__(*args, **kwargs)
+        return instance
+
+
+
 def modelcontext(model: Optional['Model']) -> Optional['Model']:
     """return the given model or try to find it in the context if there was
     none supplied.
     """
     if model is None:
         found: Optional['Model'] = Model.get_context(error_if_none=False)
         if found is None:
-            raise ValueError("No pymc3 model object on context stack.")
+            raise ValueError("No model on context stack.")
         return found
     return model
 
@@ -303,15 +348,6 @@ def logp_nojact(self):
         return logp
 
 
-class InitContextMeta(type):
-    """Metaclass that executes `__init__` of instance in it's context"""
-    def __call__(cls, *args, **kwargs):
-        instance = cls.__new__(cls, *args, **kwargs)
-        with instance:  # appends context
-            instance.__init__(*args, **kwargs)
-        return instance
-
-
 def withparent(meth):
     """Helper wrapper that passes calls to parent's instance"""
     def wrapped(self, *args, **kwargs):
@@ -566,7 +602,7 @@ def _build_joined(self, cost, args, vmap):
         return args_joined, theano.clone(cost, replace=replace)
 
 
-class Model(Context, Factor, WithMemoization, metaclass=InitContextMeta):
+class Model(Factor, WithMemoization, metaclass=ContextMeta, context_class='Model'):
     """Encapsulates the variables and likelihood factors of a model.
 
     Model class can be used for creating class based models. To create
@@ -660,7 +696,7 @@ def __new__(cls, *args, **kwargs):
         if kwargs.get('model') is not None:
             instance._parent = kwargs.get('model')
         else:
-            instance._parent = cls.get_context()
+            instance._parent = cls.get_context(error_if_none=False)
         theano_config = kwargs.get('theano_config', None)
         if theano_config is None or 'compute_test_value' not in theano_config:
             theano_config = {'compute_test_value': 'raise'}

pymc3/tests/test_data_container.py

@@ -77,7 +77,7 @@ def test_sample_after_set_data(self):
                                    atol=1e-1)
 
     def test_creation_of_data_outside_model_context(self):
-        with pytest.raises(TypeError) as error:
+        with pytest.raises((IndexError, TypeError)) as error:
             pm.Data('data', [1.1, 2.2, 3.3])
         error.match('No model on context stack')
 

pymc3/tests/test_model.py

@@ -52,10 +52,15 @@ def test_setattr_properly_works(self):
 
     def test_context_passes_vars_to_parent_model(self):
         with pm.Model() as model:
+            assert pm.model.modelcontext(None) == model
+            assert pm.Model.get_context() == model
             # a set of variables is created
-            NewModel()
+            nm = NewModel()
+            assert pm.Model.get_context() == model
             # another set of variables are created but with prefix 'another'
             usermodel2 = NewModel(name='another')
+            assert pm.Model.get_context() == model
+            assert usermodel2._parent == model
             # you can enter in a context with submodel
             with usermodel2:
                 usermodel2.Var('v3', pm.Normal.dist())

pymc3/tests/test_modelcontext.py

@@ -65,14 +65,18 @@ def test_mixed_contexts():
                     assert _DrawValuesContext.get_context() == dvcb
                     assert _DrawValuesContextBlocker.get_context() == dvcb
                 assert _DrawValuesContext.get_context() == dvc
-                assert _DrawValuesContextBlocker.get_context() is None
+                assert _DrawValuesContextBlocker.get_context() is dvc
                 assert Model.get_context() == modelB
                 assert modelcontext(None) == modelB
-            assert _DrawValuesContext.get_context() is None
+            assert _DrawValuesContext.get_context(error_if_none=False) is None
+            with raises(TypeError):
+                _DrawValuesContext.get_context()
             assert Model.get_context() == modelB
             assert modelcontext(None) == modelB
         assert Model.get_context() == modelA
         assert modelcontext(None) == modelA
-    assert Model.get_context() is None
+    assert Model.get_context(error_if_none=False) is None
+    with raises(TypeError):
+        Model.get_context(error_if_none=True)
     with raises((ValueError, TypeError)):
         modelcontext(None)