Consolidate numpy member transforms to reduce function calls

Author: correctmost

astroid/brain/brain_numpy_core_function_base.py

@@ -7,8 +7,8 @@
 import functools
 
 from astroid.brain.brain_numpy_utils import (
-    attribute_looks_like_numpy_member,
-    infer_numpy_member,
+    attribute_name_looks_like_numpy_member,
+    infer_numpy_attribute,
 )
 from astroid.inference_tip import inference_tip
 from astroid.manager import AstroidManager
@@ -25,10 +25,11 @@
 
 
 def register(manager: AstroidManager) -> None:
-    for func_name, func_src in METHODS_TO_BE_INFERRED.items():
-        inference_function = functools.partial(infer_numpy_member, func_src)
-        manager.register_transform(
-            Attribute,
-            inference_tip(inference_function),
-            functools.partial(attribute_looks_like_numpy_member, func_name),
-        )
+    manager.register_transform(
+        Attribute,
+        inference_tip(functools.partial(infer_numpy_attribute, METHODS_TO_BE_INFERRED)),
+        functools.partial(
+            attribute_name_looks_like_numpy_member,
+            frozenset(METHODS_TO_BE_INFERRED.keys()),
+        ),
+    )

astroid/brain/brain_numpy_core_multiarray.py

@@ -8,9 +8,10 @@
 
 from astroid import nodes
 from astroid.brain.brain_numpy_utils import (
-    attribute_looks_like_numpy_member,
-    infer_numpy_member,
-    name_looks_like_numpy_member,
+    attribute_name_looks_like_numpy_member,
+    infer_numpy_attribute,
+    infer_numpy_name,
+    member_name_looks_like_numpy_member,
 )
 from astroid.brain.helpers import register_module_extender
 from astroid.builder import parse
@@ -92,15 +93,15 @@ def register(manager: AstroidManager) -> None:
         manager, "numpy.core.multiarray", numpy_core_multiarray_transform
     )
 
-    for method_name, function_src in METHODS_TO_BE_INFERRED.items():
-        inference_function = functools.partial(infer_numpy_member, function_src)
-        manager.register_transform(
-            Attribute,
-            inference_tip(inference_function),
-            functools.partial(attribute_looks_like_numpy_member, method_name),
-        )
-        manager.register_transform(
-            Name,
-            inference_tip(inference_function),
-            functools.partial(name_looks_like_numpy_member, method_name),
-        )
+    method_names = frozenset(METHODS_TO_BE_INFERRED.keys())
+
+    manager.register_transform(
+        Attribute,
+        inference_tip(functools.partial(infer_numpy_attribute, METHODS_TO_BE_INFERRED)),
+        functools.partial(attribute_name_looks_like_numpy_member, method_names),
+    )
+    manager.register_transform(
+        Name,
+        inference_tip(functools.partial(infer_numpy_name, METHODS_TO_BE_INFERRED)),
+        functools.partial(member_name_looks_like_numpy_member, method_names),
+    )

astroid/brain/brain_numpy_core_numeric.py

@@ -8,8 +8,8 @@
 
 from astroid import nodes
 from astroid.brain.brain_numpy_utils import (
-    attribute_looks_like_numpy_member,
-    infer_numpy_member,
+    attribute_name_looks_like_numpy_member,
+    infer_numpy_attribute,
 )
 from astroid.brain.helpers import register_module_extender
 from astroid.builder import parse
@@ -41,10 +41,11 @@ def register(manager: AstroidManager) -> None:
         manager, "numpy.core.numeric", numpy_core_numeric_transform
     )
 
-    for method_name, function_src in METHODS_TO_BE_INFERRED.items():
-        inference_function = functools.partial(infer_numpy_member, function_src)
-        manager.register_transform(
-            Attribute,
-            inference_tip(inference_function),
-            functools.partial(attribute_looks_like_numpy_member, method_name),
-        )
+    manager.register_transform(
+        Attribute,
+        inference_tip(functools.partial(infer_numpy_attribute, METHODS_TO_BE_INFERRED)),
+        functools.partial(
+            attribute_name_looks_like_numpy_member,
+            frozenset(METHODS_TO_BE_INFERRED.keys()),
+        ),
+    )

astroid/brain/brain_numpy_utils.py

@@ -34,6 +34,21 @@ def _get_numpy_version() -> tuple[str, str, str]:
         return ("0", "0", "0")
 
 
+def infer_numpy_name(
+    sources: dict[str, str], node: Name, context: InferenceContext | None = None
+):
+    extracted_node = extract_node(sources[node.name])
+    return extracted_node.infer(context=context)
+
+
+def infer_numpy_attribute(
+    sources: dict[str, str], node: Attribute, context: InferenceContext | None = None
+):
+    extracted_node = extract_node(sources[node.attrname])
+    return extracted_node.infer(context=context)
+
+
+# TODO: Deprecate and remove this function
 def infer_numpy_member(src, node, context: InferenceContext | None = None):
     node = extract_node(src)
     return node.infer(context=context)
@@ -61,6 +76,29 @@ def _is_a_numpy_module(node: Name) -> bool:
     )
 
 
+def member_name_looks_like_numpy_member(
+    member_names: frozenset[str], node: Name
+) -> bool:
+    """
+    Returns True if the Name node's name matches a member name from numpy
+    """
+    return node.name in member_names and node.root().name.startswith("numpy")
+
+
+def attribute_name_looks_like_numpy_member(
+    member_names: frozenset[str], node: Attribute
+) -> bool:
+    """
+    Returns True if the Attribute node's name matches a member name from numpy
+    """
+    return (
+        node.attrname in member_names
+        and isinstance(node.expr, Name)
+        and _is_a_numpy_module(node.expr)
+    )
+
+
+# TODO: Deprecate and remove this function
 def name_looks_like_numpy_member(member_name: str, node: Name) -> bool:
     """
     Returns True if the Name is a member of numpy whose
@@ -69,6 +107,7 @@ def name_looks_like_numpy_member(member_name: str, node: Name) -> bool:
     return node.name == member_name and node.root().name.startswith("numpy")
 
 
+# TODO: Deprecate and remove this function
 def attribute_looks_like_numpy_member(member_name: str, node: Attribute) -> bool:
     """
     Returns True if the Attribute is a member of numpy whose