Add facility to support with-contexts

numba/__init__.py

@@ -33,6 +33,9 @@
 # Re-export jitclass
 from .jitclass import jitclass
 
+# Initialize withcontexts
+import numba.withcontexts
+
 # Keep this for backward compatibility.
 test = runtests.main
 

numba/compiler.py

@@ -429,6 +429,24 @@ def frontend_looplift(self):
                               lifted=tuple(loops), lifted_from=None)
             return cres
 
+    def frontend_withlift(self):
+        """
+        Extract with-contexts
+        """
+        main, withs = transforms.with_lifting(
+            func_ir=self.func_ir,
+            typingctx=self.typingctx,
+            targetctx=self.targetctx,
+            flags=self.flags,
+            locals=self.locals,
+            )
+        if withs:
+            cres = compile_ir(self.typingctx, self.targetctx, main,
+                              self.args, self.return_type,
+                              self.flags, self.locals,
+                              lifted=tuple(withs), lifted_from=None)
+            raise _EarlyPipelineCompletion(cres)
+
     def stage_objectmode_frontend(self):
         """
         Front-end: Analyze bytecode, generate Numba IR, infer types
@@ -449,6 +467,7 @@ def stage_nopython_frontend(self):
         """
         Type inference and legalization
         """
+        self.frontend_withlift()
         with self.fallback_context('Function "%s" failed type inference'
                                    % (self.func_id.func_name,)):
             # Type inference

numba/controlflow.py

@@ -9,7 +9,7 @@
 
 # List of bytecodes creating a new block in the control flow graph
 # (in addition to explicit jump labels).
-NEW_BLOCKERS = frozenset(['SETUP_LOOP', 'FOR_ITER'])
+NEW_BLOCKERS = frozenset(['SETUP_LOOP', 'FOR_ITER', 'SETUP_WITH'])
 
 
 class CFBlock(object):
@@ -478,6 +478,7 @@ def __init__(self, bytecode):
         self._curblock = None
         self._blockstack = []
         self._loops = []
+        self._withs = []
 
     def iterblocks(self):
         """
@@ -599,6 +600,16 @@ def op_SETUP_LOOP(self, inst):
         self.jump(inst.next)
         self._force_new_block = True
 
+    def op_SETUP_WITH(self, inst):
+        end = inst.get_jump_target()
+        self._blockstack.append(end)
+        self._withs.append((inst.offset, end))
+        # TODO: WithLifting requires the loop entry be its own block.
+        #       Forcing a new block here is the simplest solution for now.
+        #       But, we should consider other less ad-hoc ways.
+        self.jump(inst.next)
+        self._force_new_block = True
+
     def op_POP_BLOCK(self, inst):
         self._blockstack.pop()
 

numba/dataflow.py

@@ -668,6 +668,28 @@ def op_SETUP_LOOP(self, info, inst):
         self.add_syntax_block(info, LoopBlock())
         info.append(inst)
 
+    def op_SETUP_WITH(self, info, inst):
+        cm = info.pop()    # the context-manager
+        self.add_syntax_block(info, WithBlock())
+        yielded = info.make_temp()
+        info.push(yielded)
+        info.append(inst, contextmanager=cm)
+
+    def op_WITH_CLEANUP(self, info, inst):
+        """
+        Note: py2 only opcode
+        """
+        info.append(inst)
+
+    def op_WITH_CLEANUP_START(self, info, inst):
+        info.append(inst)
+
+    def op_WITH_CLEANUP_FINISH(self, info, inst):
+        info.append(inst)
+
+    def op_END_FINALLY(self, info, inst):
+        info.append(inst)
+
     def op_POP_BLOCK(self, info, inst):
         block = self.pop_syntax_block(info)
         info.append(inst)
@@ -704,7 +726,7 @@ def op_MAKE_FUNCTION(self, info, inst, MAKE_CLOSURE=False):
             if MAKE_CLOSURE:
                 closure = info.pop()
             if num_annotations > 0:
-                annotations = info.pop() 
+                annotations = info.pop()
             if num_kwdefaults > 0:
                 kwdefaults = []
                 for i in range(num_kwdefaults):
@@ -727,7 +749,7 @@ def op_MAKE_FUNCTION(self, info, inst, MAKE_CLOSURE=False):
             if inst.arg & 0x1:
                 defaults = info.pop()
         res = info.make_temp()
-        info.append(inst, name=name, code=code, closure=closure, annotations=annotations, 
+        info.append(inst, name=name, code=code, closure=closure, annotations=annotations,
                     kwdefaults=kwdefaults, defaults=defaults, res=res)
         info.push(res)
 
@@ -759,6 +781,13 @@ def __init__(self):
         self.stack_offset = None
 
 
+class WithBlock(object):
+    __slots__ = ('stack_offset',)
+
+    def __init__(self):
+        self.stack_offset = None
+
+
 class BlockInfo(object):
     def __init__(self, block, offset, incoming_blocks):
         self.block = block

numba/datamodel/models.py

@@ -288,6 +288,7 @@ def __init__(self, dmm, fe_type):
 @register_default(types.Object)
 @register_default(types.Module)
 @register_default(types.Phantom)
+@register_default(types.ContextManager)
 @register_default(types.Dispatcher)
 @register_default(types.ExceptionClass)
 @register_default(types.Dummy)

numba/dispatcher.py

@@ -10,14 +10,13 @@
 import uuid
 import weakref
 
-import numba
 from numba import _dispatcher, compiler, utils, types, config, errors
 from numba.typeconv.rules import default_type_manager
 from numba import sigutils, serialize, typing
 from numba.typing.templates import fold_arguments
-from numba.typing.typeof import Purpose, typeof, typeof_impl
+from numba.typing.typeof import Purpose, typeof
 from numba.bytecode import get_code_object
-from numba.six import create_bound_method, next, reraise
+from numba.six import create_bound_method, reraise
 from .caching import NullCache, FunctionCache
 
 
@@ -681,9 +680,9 @@ def stats(self):
             )
 
 
-class LiftedLoop(_DispatcherBase):
+class LiftedCode(_DispatcherBase):
     """
-    Implementation of the hidden dispatcher objects used for lifted loop
+    Implementation of the hidden dispatcher objects used for lifted code
     (a lifted loop is really compiled as a separate function).
     """
     _fold_args = False
@@ -707,6 +706,11 @@ def get_source_location(self):
         """
         return self.func_ir.loc.line
 
+    def _pre_compile(self, args, return_type, flags):
+        """Pre-compile actions
+        """
+        pass
+
     def compile(self, sig):
         # Use cache and compiler in a critical section
         with compiler.lock_compiler:
@@ -722,7 +726,8 @@ def compile(self, sig):
                 if existing is not None:
                     return existing.entry_point
 
-                assert not flags.enable_looplift, "Enable looplift flags is on"
+                self._pre_compile(args, return_type, flags)
+
                 # Clone IR to avoid mutation in rewrite pass
                 cloned_func_ir = self.func_ir.copy()
                 cres = compiler.compile_ir(typingctx=self.typingctx,
@@ -741,6 +746,38 @@ def compile(self, sig):
                 return cres.entry_point
 
 
+class LiftedLoop(LiftedCode):
+    def _pre_compile(self, args, return_type, flags):
+        assert not flags.enable_looplift, "Enable looplift flags is on"
+
+
+class LiftedWith(LiftedCode):
+    @property
+    def _numba_type_(self):
+        return types.Dispatcher(self)
+
+    def get_call_template(self, args, kws):
+        """
+        Get a typing.ConcreteTemplate for this dispatcher and the given
+        *args* and *kws* types.  This enables the resolving of the return type.
+
+        A (template, pysig, args, kws) tuple is returned.
+        """
+        # Ensure an overload is available
+        if self._can_compile:
+            self.compile(tuple(args))
+
+        pysig = None
+        # Create function type for typing
+        func_name = self.py_func.__name__
+        name = "CallTemplate({0})".format(func_name)
+        # The `key` isn't really used except for diagnosis here,
+        # so avoid keeping a reference to `cfunc`.
+        call_template = typing.make_concrete_template(
+            name, key=func_name, signatures=self.nopython_signatures)
+        return call_template, pysig, args, kws
+
+
 # Initialize typeof machinery
 _dispatcher.typeof_init(
     OmittedArg,

numba/interpreter.py

@@ -217,7 +217,7 @@ def code_names(self):
 
     @property
     def code_cellvars(self):
-        return self.bytecode.co_cellvars 
+        return self.bytecode.co_cellvars
 
     @property
     def code_freevars(self):
@@ -272,7 +272,7 @@ def get(self, name):
         # See Parameter class in inspect.py (from Python source)
         if name[0] == '.' and name[1:].isdigit():
             name = 'implicit{}'.format(name[1:])
- 
+
         # Try to simplify the variable lookup by returning an earlier
         # variable assigned to *name*.
         var = self.assigner.get_assignment_source(name)
@@ -612,6 +612,28 @@ def op_SETUP_LOOP(self, inst):
         loop = ir.Loop(inst.offset, exit=(inst.next + inst.arg))
         self.syntax_blocks.append(loop)
 
+    def op_SETUP_WITH(self, inst, contextmanager):
+        assert self.blocks[inst.offset] is self.current_block
+        exitpt = inst.next + inst.arg
+        wth = ir.With(inst.offset, exit=exitpt)
+        self.syntax_blocks.append(wth)
+        self.current_block.append(ir.EnterWith(
+            contextmanager=self.get(contextmanager),
+            begin=inst.offset, end=exitpt, loc=self.loc,
+            ))
+
+    def op_WITH_CLEANUP(self, inst):
+        "no-op"
+
+    def op_WITH_CLEANUP_START(self, inst):
+        "no-op"
+
+    def op_WITH_CLEANUP_FINISH(self, inst):
+        "no-op"
+
+    def op_END_FINALLY(self, inst):
+        "no-op"
+
     if PYVERSION < (3, 6):
 
         def op_CALL_FUNCTION(self, inst, func, args, kws, res, vararg):
@@ -969,7 +991,7 @@ def op_MAKE_FUNCTION(self, inst, name, code, closure, annotations, kwdefaults, d
 
     def op_MAKE_CLOSURE(self, inst, name, code, closure, annotations, kwdefaults, defaults, res):
         self.op_MAKE_FUNCTION(inst, name, code, closure, annotations, kwdefaults, defaults, res)
-    
+
     def op_LOAD_CLOSURE(self, inst, res):
         n_cellvars = len(self.code_cellvars)
         if inst.arg < n_cellvars:

numba/ir.py

@@ -570,6 +570,31 @@ def list_vars(self):
         return [self.value]
 
 
+class EnterWith(Stmt):
+    """Enter a "with" context
+    """
+    def __init__(self, contextmanager, begin, end, loc):
+        """
+        Parameters
+        ----------
+        contextmanager : IR value
+        begin, end : int
+            The beginning and the ending offset of the with-body.
+        loc : int
+            Source location
+        """
+        self.contextmanager = contextmanager
+        self.begin = begin
+        self.end = end
+        self.loc = loc
+
+    def __str__(self):
+        return 'enter_with {}'.format(self.contextmanager)
+
+    def list_vars(self):
+        return [self.contextmanager]
+
+
 class Arg(object):
     def __init__(self, name, index, loc):
         self.name = name
@@ -880,6 +905,8 @@ def __repr__(self):
 
 
 class Loop(object):
+    """Describes a loop-block
+    """
     __slots__ = "entry", "exit"
 
     def __init__(self, entry, exit):
@@ -891,6 +918,20 @@ def __repr__(self):
         return "Loop(entry=%s, exit=%s)" % args
 
 
+class With(object):
+    """Describes a with-block
+    """
+    __slots__ = "entry", "exit"
+
+    def __init__(self, entry, exit):
+        self.entry = entry
+        self.exit = exit
+
+    def __repr__(self):
+        args = self.entry, self.exit
+        return "With(entry=%s, exit=%s)" % args
+
+
 class FunctionIR(object):
 
     def __init__(self, blocks, is_generator, func_id, loc,
@@ -1012,4 +1053,8 @@ def dump_generator_info(self, file=None):
 
 
 # A stub for undefined global reference
-UNDEFINED = object()
+class UndefinedType(object):
+    def __repr__(self):
+        return "Undefined"
+
+UNDEFINED = UndefinedType()