diff --git a/ChangeLog b/ChangeLog index b612aee083..54826817ea 100644 --- a/ChangeLog +++ b/ChangeLog @@ -6,6 +6,11 @@ What's New in astroid 2.4.0? ============================ Release Date: TBA +* Added a call to ``register_transform`` for all functions of the ``brain_numpy_core_multiarray`` + module in case the current node is an instance of ``astroid.Name`` + + Close #666 + * Added some functions of the ``numpy.core.multiarray`` module Close PyCQA/pylint#3208 diff --git a/astroid/brain/brain_numpy_core_multiarray.py b/astroid/brain/brain_numpy_core_multiarray.py index 182184366e..8f5e32cc02 100644 --- a/astroid/brain/brain_numpy_core_multiarray.py +++ b/astroid/brain/brain_numpy_core_multiarray.py @@ -85,3 +85,8 @@ def vdot(a, b): astroid.inference_tip(inference_function), functools.partial(looks_like_numpy_member, method_name), ) + astroid.MANAGER.register_transform( + astroid.Name, + astroid.inference_tip(inference_function), + functools.partial(looks_like_numpy_member, method_name), + ) diff --git a/astroid/brain/brain_numpy_utils.py b/astroid/brain/brain_numpy_utils.py index 2bad01eed4..c51922c38d 100644 --- a/astroid/brain/brain_numpy_utils.py +++ b/astroid/brain/brain_numpy_utils.py @@ -48,9 +48,17 @@ def looks_like_numpy_member( :param node: node to test :return: True if the node is a member of numpy """ - return ( + if ( isinstance(node, astroid.Attribute) and node.attrname == member_name and isinstance(node.expr, astroid.Name) and _is_a_numpy_module(node.expr) - ) + ): + return True + if ( + isinstance(node, astroid.Name) + and node.name == member_name + and node.root().name.startswith("numpy") + ): + return True + return False