diff --git a/astroid/brain/brain_numpy_core_function_base.py b/astroid/brain/brain_numpy_core_function_base.py index bd218efa57..f69826d55e 100644 --- a/astroid/brain/brain_numpy_core_function_base.py +++ b/astroid/brain/brain_numpy_core_function_base.py @@ -6,7 +6,10 @@ import functools -from astroid.brain.brain_numpy_utils import infer_numpy_member, looks_like_numpy_member +from astroid.brain.brain_numpy_utils import ( + attribute_looks_like_numpy_member, + infer_numpy_member, +) from astroid.inference_tip import inference_tip from astroid.manager import AstroidManager from astroid.nodes.node_classes import Attribute @@ -25,5 +28,5 @@ AstroidManager().register_transform( Attribute, inference_tip(inference_function), - functools.partial(looks_like_numpy_member, func_name), + functools.partial(attribute_looks_like_numpy_member, func_name), ) diff --git a/astroid/brain/brain_numpy_core_multiarray.py b/astroid/brain/brain_numpy_core_multiarray.py index 4b2fe63c02..e9c7bacfce 100644 --- a/astroid/brain/brain_numpy_core_multiarray.py +++ b/astroid/brain/brain_numpy_core_multiarray.py @@ -6,7 +6,11 @@ import functools -from astroid.brain.brain_numpy_utils import infer_numpy_member, looks_like_numpy_member +from astroid.brain.brain_numpy_utils import ( + attribute_looks_like_numpy_member, + infer_numpy_member, + name_looks_like_numpy_member, +) from astroid.brain.helpers import register_module_extender from astroid.builder import parse from astroid.inference_tip import inference_tip @@ -91,10 +95,10 @@ def vdot(a, b): AstroidManager().register_transform( Attribute, inference_tip(inference_function), - functools.partial(looks_like_numpy_member, method_name), + functools.partial(attribute_looks_like_numpy_member, method_name), ) AstroidManager().register_transform( Name, inference_tip(inference_function), - functools.partial(looks_like_numpy_member, method_name), + functools.partial(name_looks_like_numpy_member, method_name), ) diff --git a/astroid/brain/brain_numpy_core_numeric.py b/astroid/brain/brain_numpy_core_numeric.py index c5c816f00b..6fd23a857f 100644 --- a/astroid/brain/brain_numpy_core_numeric.py +++ b/astroid/brain/brain_numpy_core_numeric.py @@ -6,7 +6,10 @@ import functools -from astroid.brain.brain_numpy_utils import infer_numpy_member, looks_like_numpy_member +from astroid.brain.brain_numpy_utils import ( + attribute_looks_like_numpy_member, + infer_numpy_member, +) from astroid.brain.helpers import register_module_extender from astroid.builder import parse from astroid.inference_tip import inference_tip @@ -42,5 +45,5 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None): ret AstroidManager().register_transform( Attribute, inference_tip(inference_function), - functools.partial(looks_like_numpy_member, method_name), + functools.partial(attribute_looks_like_numpy_member, method_name), ) diff --git a/astroid/brain/brain_numpy_utils.py b/astroid/brain/brain_numpy_utils.py index 5867b6f9c2..47f24433bd 100644 --- a/astroid/brain/brain_numpy_utils.py +++ b/astroid/brain/brain_numpy_utils.py @@ -8,7 +8,7 @@ from astroid.builder import extract_node from astroid.context import InferenceContext -from astroid.nodes.node_classes import Attribute, Import, Name, NodeNG +from astroid.nodes.node_classes import Attribute, Import, Name # Class subscript is available in numpy starting with version 1.20.0 NUMPY_VERSION_TYPE_HINTS_SUPPORT = ("1", "20", "0") @@ -61,26 +61,21 @@ def _is_a_numpy_module(node: Name) -> bool: ) -def looks_like_numpy_member(member_name: str, node: NodeNG) -> bool: +def name_looks_like_numpy_member(member_name: str, node: Name) -> bool: """ - Returns True if the node is a member of numpy whose + Returns True if the Name is a member of numpy whose name is member_name. + """ + return node.name == member_name and node.root().name.startswith("numpy") - :param member_name: name of the member - :param node: node to test - :return: True if the node is a member of numpy + +def attribute_looks_like_numpy_member(member_name: str, node: Attribute) -> bool: """ - if ( - isinstance(node, Attribute) - and node.attrname == member_name + Returns True if the Attribute is a member of numpy whose + name is member_name. + """ + return ( + node.attrname == member_name and isinstance(node.expr, Name) and _is_a_numpy_module(node.expr) - ): - return True - if ( - isinstance(node, Name) - and node.name == member_name - and node.root().name.startswith("numpy") - ): - return True - return False + )