Repository: numba/numba

Issue:
Optional automatic JIT of called functions
(Based on discussion in #2630)

A good use case for Numba is transparent compilation of user-defined functions as part of a larger API.  This works if the user supplies a monolithic UDF that only calls functions Numba is already aware of, then you just wrap the UDF with `@jit`.  If the user has created helper functions that are called from the UDF, the only possible compilation mode for the main UDF would be object mode or do a lot of hackery to rewrite the UDF to call JIT compiled versions of the helper functions that are created on behalf of the user.

This would be vastly simplified if one could have Numba transform the function body to replace all calls to functions (not already known to the compiler) with JIT-wrapped versions of those functions.  For lack of a better name, this option to `@jit` could be indicated with `callees=True`.

Some design thoughts for discussion:
* I think `callees` should default to False: Compiling undecorated functions is unexpected behavior we only want to attempt if the user opts in.  Existing code would be unaffected by this change.
* The compilation of the callees should inherit the JIT options of the caller (nopython, error_mode, parallel, fastmath, etc).  This would include `callees=True` itself, which automatically makes the compilation recursive.
* To avoid accidentally trying to compile all of matplotlib (for example), we need to ensure that this can fail fast:
  - Make sure to do the recursive compilation depth first, so we can abort before consuming a ton of memory?
  - Should we only allow `callees=True` in conjunction with nopython mode?  This is going to be the primary use case (object mode is not going to do much at the moment), and being more restrictive, nopython mode is more likely to cut off a futile compilation attempt early. 
  - Should we have a max recursion depth for compilation?


Discussion:
Bikeshedding on the argument name, `recurse=True` seems more obvious to me. 

...or `recurse_depth=0` as a default and then the user could override that to allow recursive compilation of nested functions to any arbitrary depth.
Yeah, I initially avoided that name because I was afraid it would be confused somehow with support recursively callable functions, but in retrospect, that was silly on my part.

+1 for using `recurse`.   recurse_depth (if we do it at all) should be a separate option that hopefully most people don't ever need to modify.

Resolving patch:
diff --git a/numba/__init__.py b/numba/__init__.py
--- a/numba/__init__.py
+++ b/numba/__init__.py
@@ -26,7 +26,8 @@
 from .smartarray import SmartArray
 
 # Re-export decorators
-from .decorators import autojit, cfunc, generated_jit, jit, njit, stencil
+from .decorators import (autojit, cfunc, generated_jit, jit, njit, stencil,
+                         jit_module)
 
 # Re-export vectorize decorators and the thread layer querying function
 from .npyufunc import vectorize, guvectorize, threading_layer
@@ -57,6 +58,7 @@
     jitclass
     njit
     stencil
+    jit_module
     typeof
     prange
     gdb
diff --git a/numba/decorators.py b/numba/decorators.py
--- a/numba/decorators.py
+++ b/numba/decorators.py
@@ -6,6 +6,8 @@
 
 import sys
 import warnings
+import inspect
+import logging
 
 from . import config, sigutils
 from .errors import DeprecationError, NumbaDeprecationWarning
@@ -13,6 +15,8 @@
 from .stencil import stencil
 
 
+_logger = logging.getLogger(__name__)
+
 
 # -----------------------------------------------------------------------------
 # Decorators
@@ -246,3 +250,27 @@ def wrapper(func):
         return res
 
     return wrapper
+
+
+def jit_module(**kwargs):
+    """ Automatically ``jit``-wraps functions defined in a Python module
+
+    Note that ``jit_module`` should only be called at the end of the module to
+    be jitted. In addition, only functions which are defined in the module
+    ``jit_module`` is called from are considered for automatic jit-wrapping.
+    See the Numba documentation for more information about what can/cannot be
+    jitted.
+
+    :param kwargs: Keyword arguments to pass to ``jit`` such as ``nopython``
+                   or ``error_model``.
+
+    """
+    # Get the module jit_module is being called from
+    frame = inspect.stack()[1]
+    module = inspect.getmodule(frame[0])
+    # Replace functions in module with jit-wrapped versions
+    for name, obj in module.__dict__.items():
+        if inspect.isfunction(obj) and inspect.getmodule(obj) == module:
+            _logger.debug("Auto decorating function {} from module {} with jit "
+                          "and options: {}".format(obj, module.__name__, kwargs))
+            module.__dict__[name] = jit(obj, **kwargs)


Test patch:
diff --git a/numba/tests/test_jit_module.py b/numba/tests/test_jit_module.py
new file mode 100644
--- /dev/null
+++ b/numba/tests/test_jit_module.py
@@ -0,0 +1,173 @@
+
+import os
+import sys
+import shutil
+import inspect
+import importlib
+import contextlib
+import uuid
+import numpy as np
+import logging
+
+import numba.unittest_support as unittest
+from numba import dispatcher
+from numba.utils import StringIO
+from numba.tests.support import temp_directory, SerialMixin
+
+
+@contextlib.contextmanager
+def captured_logs(l):
+    try:
+        buffer = StringIO()
+        handler = logging.StreamHandler(buffer)
+        l.addHandler(handler)
+        yield buffer
+    finally:
+        l.removeHandler(handler)
+
+
+class TestJitModule(SerialMixin, unittest.TestCase):
+
+    source_lines = """
+from numba import jit_module
+
+def inc(x):
+    return x + 1
+
+def add(x, y):
+    return x + y
+
+def inc_add(x):
+    y = inc(x)
+    return add(x, y)
+
+import numpy as np
+mean = np.mean
+
+class Foo(object):
+    pass
+
+jit_module({jit_options})
+"""
+
+    def _format_jit_options(self, **jit_options):
+        if not jit_options:
+            return ''
+        out = []
+        for key, value in jit_options.items():
+            if isinstance(value, str):
+                value = '"{}"'.format(value)
+            out.append('{}={}'.format(key, value))
+        return ', '.join(out)
+
+    @contextlib.contextmanager
+    def create_temp_jitted_module(self, source_lines=None, **jit_options):
+        # Use try/finally so cleanup happens even when an exception is raised
+        try:
+            if source_lines is None:
+                source_lines = self.source_lines
+            tempdir = temp_directory('test_jit_module')
+            # Generate random module name
+            temp_module_name = 'test_module_' + str(uuid.uuid4()).replace('-', '_')
+            temp_module_path = os.path.join(tempdir, temp_module_name + '.py')
+
+            jit_options = self._format_jit_options(**jit_options)
+            with open(temp_module_path, 'w') as f:
+                lines = source_lines.format(jit_options=jit_options)
+                f.write(lines)
+            # Add test_module to sys.path so it can be imported
+            sys.path.insert(0, tempdir)
+            test_module = importlib.import_module(temp_module_name)
+            yield test_module
+        finally:
+            sys.modules.pop(temp_module_name, None)
+            sys.path.remove(tempdir)
+            shutil.rmtree(tempdir)
+
+    def test_create_temp_jitted_module(self):
+        sys_path_original = list(sys.path)
+        sys_modules_original = dict(sys.modules)
+        with self.create_temp_jitted_module() as test_module:
+            temp_module_dir = os.path.dirname(test_module.__file__)
+            self.assertEqual(temp_module_dir, sys.path[0])
+            self.assertEqual(sys.path[1:], sys_path_original)
+            self.assertTrue(test_module.__name__ in sys.modules)
+        # Test that modifications to sys.path / sys.modules are reverted
+        self.assertEqual(sys.path, sys_path_original)
+        self.assertEqual(sys.modules, sys_modules_original)
+
+    def test_create_temp_jitted_module_with_exception(self):
+        try:
+            sys_path_original = list(sys.path)
+            sys_modules_original = dict(sys.modules)
+            with self.create_temp_jitted_module():
+                raise ValueError("Something went wrong!")
+        except ValueError:
+            # Test that modifications to sys.path / sys.modules are reverted
+            self.assertEqual(sys.path, sys_path_original)
+            self.assertEqual(sys.modules, sys_modules_original)
+
+    def test_jit_module(self):
+        with self.create_temp_jitted_module() as test_module:
+            self.assertIsInstance(test_module.inc, dispatcher.Dispatcher)
+            self.assertIsInstance(test_module.add, dispatcher.Dispatcher)
+            self.assertIsInstance(test_module.inc_add, dispatcher.Dispatcher)
+            self.assertTrue(test_module.mean is np.mean)
+            self.assertTrue(inspect.isclass(test_module.Foo))
+
+            # Test output of jitted functions is as expected
+            x, y = 1.7, 2.3
+            self.assertEqual(test_module.inc(x), test_module.inc.py_func(x))
+            self.assertEqual(test_module.add(x, y), test_module.add.py_func(x, y))
+            self.assertEqual(test_module.inc_add(x), test_module.inc_add.py_func(x))
+
+    def test_jit_module_jit_options(self):
+        jit_options = {"nopython": True,
+                       "nogil": False,
+                       "error_model": "numpy",
+                       }
+        with self.create_temp_jitted_module(**jit_options) as test_module:
+            self.assertEqual(test_module.inc.targetoptions, jit_options)
+
+    def test_jit_module_jit_options_override(self):
+        source_lines = """
+from numba import jit, jit_module
+
+@jit(nogil=True, forceobj=True)
+def inc(x):
+    return x + 1
+
+def add(x, y):
+    return x + y
+
+jit_module({jit_options})
+"""
+        jit_options = {"nopython": True,
+                       "error_model": "numpy",
+                       }
+        with self.create_temp_jitted_module(source_lines=source_lines, **jit_options) as test_module:
+            self.assertEqual(test_module.add.targetoptions, jit_options)
+            # Test that manual jit-wrapping overrides jit_module options
+            self.assertEqual(test_module.inc.targetoptions, {'nogil': True, 'forceobj': True})
+
+    def test_jit_module_logging_output(self):
+        logger = logging.getLogger('numba.decorators')
+        logger.setLevel(logging.DEBUG)
+        jit_options = {"nopython": True,
+                       "error_model": "numpy",
+                       }
+        with captured_logs(logger) as logs:
+            with self.create_temp_jitted_module(**jit_options) as test_module:
+                logs = logs.getvalue()
+                expected = ["Auto decorating function",
+                            "from module {}".format(test_module.__name__),
+                            "with jit and options: {}".format(jit_options)]
+                self.assertTrue(all(i in logs for i in expected))
+
+    def test_jit_module_logging_level(self):
+        logger = logging.getLogger('numba.decorators')
+        # Test there's no logging for INFO level
+        logger.setLevel(logging.INFO)
+        with captured_logs(logger) as logs:
+            with self.create_temp_jitted_module():
+                self.assertEqual(logs.getvalue(), '')