Merge pull request #4064 from juliantaylor/import-speed

ENH: improve add_newdocs performance

2 files changed, 14 insertions, 5 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index a70f74f60..9176d9950 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -3110,15 +3110,14 @@ def add_newdoc(place, obj, doc):
     that the docstrings were changed.
        """
     try:
-        new = {}
-        exec('from %s import %s' % (place, obj), new)
+        new = getattr(__import__(place, globals(), {}, [obj]), obj)
         if isinstance(doc, str):
-            add_docstring(new[obj], doc.strip())
+            add_docstring(new, doc.strip())
         elif isinstance(doc, tuple):
-            add_docstring(getattr(new[obj], doc[0]), doc[1].strip())
+            add_docstring(getattr(new, doc[0]), doc[1].strip())
         elif isinstance(doc, list):
             for val in doc:
-                add_docstring(getattr(new[obj], val[0]), val[1].strip())
+                add_docstring(getattr(new, val[0]), val[1].strip())
     except:
         pass
 
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index f91ab8aa1..6829e06ae 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1808,6 +1808,7 @@ class TestMedian(TestCase):
         a = MySubClass([1,2,3])
         assert_equal(np.median(a), -7)
 
+
 class TestAdd_newdoc_ufunc(TestCase):
 
     def test_ufunc_arg(self):
@@ -1818,5 +1819,14 @@ class TestAdd_newdoc_ufunc(TestCase):
         assert_raises(TypeError, add_newdoc_ufunc, np.add, 3)
 
 
+class TestAdd_newdoc(TestCase):
+    def test_add_doc(self):
+        # test np.add_newdoc
+        tgt = "Current flat index into the array."
+        self.assertEqual(np.core.flatiter.index.__doc__[:len(tgt)], tgt)
+        self.assertTrue(len(np.core.ufunc.identity.__doc__) > 300)
+        self.assertTrue(len(np.lib.index_tricks.mgrid.__doc__) > 300)
+
+
 if __name__ == "__main__":
     run_module_suite()