withmagic.py

# -*- coding: utf-8 -*-
"""
%with and %endwith magics.
"""
#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import re

from IPython.core.error import UsageError
from IPython.core.magic import Magics, magics_class, line_magic

#-----------------------------------------------------------------------------
# Functions and classes
#-----------------------------------------------------------------------------

@magics_class
class WithMagics(Magics):
    """Provides the %with and %endwith magics."""

    def __init__(self, shell):
        super(WithMagics, self).__init__(shell)
        self._re_with = re.compile(r'^(?P<code>.*)\s+as\s+(?P<name>[_A-Za-z][_A-Za-z0-9]*):$')
        self._contexts = []

    @line_magic('with')
    def with_(self, line=''):
        m = self._re_with.match(line)
        if not m:
            raise UsageError("Error parsing statement.")

        code = m.group('code')
        name = m.group('name')

        shell = self.shell
        context = eval(code, shell.user_global_ns, shell.user_ns)

        obj = context.__enter__()
        print("Context manager %s started" % name)

        # Return value to the user
        shell.user_ns[name] = obj

        self._contexts.append((name, context))

    @line_magic
    def endwith(self, line=''):
        name, context = self._contexts.pop()
        print("Context manager %s ended" % name)
        context.__exit__(None, None, None)

def load_ipython_extension(ip):
    """Load the extension in IPython."""
    ip.register_magics(WithMagics)