diff options
| -rw-r--r-- | astroid/brain/pynumpy.py | 20 | ||||
| -rw-r--r-- | astroid/tests/unittest_brain.py | 2 |
2 files changed, 19 insertions, 3 deletions
diff --git a/astroid/brain/pynumpy.py b/astroid/brain/pynumpy.py index cb43339..75f4f18 100644 --- a/astroid/brain/pynumpy.py +++ b/astroid/brain/pynumpy.py @@ -24,7 +24,7 @@ import astroid # TODO(cpopa): drop when understanding augmented assignments def numpy_core_transform(): - module = astroid.parse(''' + return astroid.parse(''' from numpy.core import numeric from numpy.core import fromnumeric from numpy.core import defchararray @@ -41,6 +41,22 @@ def numpy_core_transform(): getlimits.__all__ + shape_base.__all__) ''') - return module + + +def numpy_transform(): + return astroid.parse(''' + from numpy import core + from numpy import matrixlib as _mat + from numpy import lib + __all__ = ['add_newdocs', + 'ModuleDeprecationWarning', + 'VisibleDeprecationWarning', 'linalg', 'fft', 'random', + 'ctypeslib', 'ma', + '__version__', 'pkgload', 'PackageLoader', + 'show_config'] + core.__all__ + _mat.__all__ + lib.__all__ + + ''') + astroid.register_module_extender(astroid.MANAGER, 'numpy.core', numpy_core_transform) +astroid.register_module_extender(astroid.MANAGER, 'numpy', numpy_transform) diff --git a/astroid/tests/unittest_brain.py b/astroid/tests/unittest_brain.py index 510b4c0..02c405b 100644 --- a/astroid/tests/unittest_brain.py +++ b/astroid/tests/unittest_brain.py @@ -435,7 +435,7 @@ class NumpyBrainTest(unittest.TestCase): numpy.ones #@ ''') inferred = next(node.infer()) - self.assertIsInstance(inferred, nodes.Function) + self.assertIsInstance(inferred, nodes.FunctionDef) if __name__ == '__main__': |