Merge pull request #13529 from shoyer/array-function-exec

Use exec() instead array_function_dispatch to improve tracebacks

1 files changed, 22 insertions, 2 deletions

diff --git a/numpy/core/overrides.py b/numpy/core/overrides.py
index 9f91adc83..bb6d8694a 100644
--- a/numpy/core/overrides.py
+++ b/numpy/core/overrides.py
@@ -1,6 +1,7 @@
 """Implementation of __array_function__ overrides from NEP-18."""
 import collections
 import functools
+import textwrap
 
 from numpy.core._multiarray_umath import (
     add_docstring, implement_array_function, _get_implementing_args)
@@ -143,11 +144,30 @@ def array_function_dispatch(dispatcher, module=None, verify=True,
         if docs_from_dispatcher:
             add_docstring(implementation, dispatcher.__doc__)
 
+        # Equivalently, we could define this function directly instead of using
+        # exec. This version has the advantage of giving the helper function a
+        # more interpettable name. Otherwise, the original function does not
+        # show up at all in many cases, e.g., if it's written in C or if the
+        # dispatcher gets an invalid keyword argument.
+        source = textwrap.dedent("""
         @functools.wraps(implementation)
-        def public_api(*args, **kwargs):
+        def {name}(*args, **kwargs):
             relevant_args = dispatcher(*args, **kwargs)
             return implement_array_function(
-                implementation, public_api, relevant_args, args, kwargs)
+                implementation, {name}, relevant_args, args, kwargs)
+        """).format(name=implementation.__name__)
+
+        source_object = compile(
+            source, filename='<__array_function__ internals>', mode='exec')
+        scope = {
+            'implementation': implementation,
+            'dispatcher': dispatcher,
+            'functools': functools,
+            'implement_array_function': implement_array_function,
+        }
+        exec(source_object, scope)
+
+        public_api = scope[implementation.__name__]
 
         if module is not None:
             public_api.__module__ = module