Repository: numba/numba

Issue:
Objmode with lifting not using custom pipeline
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I am using the latest released version of Numba (most recent is visible in
 the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [ ] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->
Object mode `with` lifting stops the current compiler pipeline and calls the compiler again, but uses the base pipeline class instead of the user pipeline class.


Discussion:

Resolving patch:
diff --git a/numba/compiler.py b/numba/compiler.py
--- a/numba/compiler.py
+++ b/numba/compiler.py
@@ -443,7 +443,8 @@ def stage_frontend_withlift(self):
             cres = compile_ir(self.typingctx, self.targetctx, main,
                               self.args, self.return_type,
                               self.flags, self.locals,
-                              lifted=tuple(withs), lifted_from=None)
+                              lifted=tuple(withs), lifted_from=None,
+                              pipeline_class=type(self))
             raise _EarlyPipelineCompletion(cres)
 
     def stage_objectmode_frontend(self):
@@ -905,15 +906,16 @@ def compile_extra(typingctx, targetctx, func, args, return_type, flags,
 
 
 def compile_ir(typingctx, targetctx, func_ir, args, return_type, flags,
-               locals, lifted=(), lifted_from=None, library=None):
+               locals, lifted=(), lifted_from=None, library=None,
+               pipeline_class=Pipeline):
     """
     Compile a function with the given IR.
 
     For internal use only.
     """
 
-    pipeline = Pipeline(typingctx, targetctx, library,
-                        args, return_type, flags, locals)
+    pipeline = pipeline_class(typingctx, targetctx, library,
+                              args, return_type, flags, locals)
     return pipeline.compile_ir(func_ir=func_ir, lifted=lifted,
                                lifted_from=lifted_from)
 


Test patch:
diff --git a/numba/tests/test_pipeline.py b/numba/tests/test_pipeline.py
--- a/numba/tests/test_pipeline.py
+++ b/numba/tests/test_pipeline.py
@@ -1,7 +1,8 @@
 from __future__ import print_function
 
 from numba.compiler import Pipeline
-from numba import jit, generated_jit, types
+from numba import jit, generated_jit, types, objmode
+from numba.ir import FunctionIR
 from .support import TestCase
 
 
@@ -18,6 +19,12 @@ def compile_extra(self, func):
                 self.custom_pipeline_cache.append(func)
                 return super(CustomPipeline, self).compile_extra(func)
 
+            def compile_ir(self, func_ir, *args, **kwargs):
+                # Store the compiled function
+                self.custom_pipeline_cache.append(func_ir)
+                return super(CustomPipeline, self).compile_ir(
+                    func_ir, *args, **kwargs)
+
         self.pipeline_class = CustomPipeline
 
     def test_jit_custom_pipeline(self):
@@ -45,3 +52,24 @@ def foo(x):
         self.assertEqual(foo(5), 5)
         self.assertListEqual(self.pipeline_class.custom_pipeline_cache,
                              [inner])
+
+    def test_objmode_custom_pipeline(self):
+        self.assertListEqual(self.pipeline_class.custom_pipeline_cache, [])
+
+        @jit(pipeline_class=self.pipeline_class)
+        def foo(x):
+            with objmode(x="intp"):
+                x += int(0x1)
+            return x
+
+        arg = 123
+        self.assertEqual(foo(arg), arg + 1)
+        # Two items in the list.
+        self.assertEqual(len(self.pipeline_class.custom_pipeline_cache), 2)
+        # First item is the `foo` function
+        first = self.pipeline_class.custom_pipeline_cache[0]
+        self.assertIs(first, foo.py_func)
+        # Second item is a FunctionIR of the obj-lifted function
+        second = self.pipeline_class.custom_pipeline_cache[1]
+        self.assertIsInstance(second, FunctionIR)
+