Python for Everyone!
Oregon Curriculum Network

Spooky Castle

Context Managers in Python

Halloween is nigh. Or lets pretend it is.

Lets explore context manager syntax in the context of Halloween...


In [4]:
from random import choice

class Trick(Exception):
    def __init__(self):
        self.value = "Goblins & Ghosts!"
    
class Halloween:
    
    def __init__(self, arg=None):
        self.testing = arg
    
    def __enter__(self):
        self.where = "Spooky Castle"
        print("Welcome...")
        self.trick_or_treat = ["Trick", "Treat"]
        self.candy = [ ]
        return self
    
    def __exit__(self, *uh_oh):  # catch any exception info
        if uh_oh[0]: 
            print("Trick!")
            print(uh_oh[1].value)  # lets look inside the exception
            return False
        return True

try:
    with Halloween("Testing 1-2-3") as obj:
        print(obj.testing)
        if choice(obj.trick_or_treat) == "Trick":
            raise Trick    
except Trick:
    print("Exception raised!")


Welcome...
Testing 1-2-3

Notice that test_trick below defines a function internally, and calls it, expecting an Exception to be raised. Defining functions internally to methods for testing purposes is a nice way to avoid creating any test functions globally. No need to clean up. No worries about persistence of "hidden state" across tests (tests should not depend on one another).

One needs some fancy footwork to run the unittests in a Jupyter Notebook, but it's still very doable and fun!


In [5]:
import unittest

class TestCastle(unittest.TestCase):
    
    def test_candy(self):
        outer = ""
        with Halloween() as context:
            outer = context.candy
        self.assertEqual(outer, [], "Not OK!")
        
    def test_trick(self):
        outer = ""
        def func():
            with Halloween() as context:
                raise Trick
        self.assertRaises(Trick, func)
        
a = TestCastle()  # the test suite
suite = unittest.TestLoader().loadTestsFromModule(a) # fancy boilerplate
unittest.TextTestRunner().run(suite)  # run the test suite


..
Welcome...
Welcome...
Trick!
Goblins & Ghosts!
----------------------------------------------------------------------
Ran 2 tests in 0.006s

OK
Out[5]:
<unittest.runner.TextTestResult run=2 errors=0 failures=0>

In [ ]: