Add brain for numpy core module einsumfunc. (#1656)

Co-authored-by: Daniël van Noord <[email protected]>

Author: mbyrnepr2

ChangeLog

@@ -95,6 +95,10 @@ Release date: TBA
 * Fix test for Python ``3.11``. In some instances ``err.__traceback__`` will
   be uninferable now.
 
+* Add brain for numpy core module ``einsumfunc``.
+
+  Closes PyCQA/pylint#5821
+
 * Infer the ``DictUnpack`` value for ``Dict.getitem`` calls.
 
   Closes #1195

astroid/brain/brain_numpy_core_einsumfunc.py

@@ -0,0 +1,27 @@
+# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
+# For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
+# Copyright (c) https://github.com/PyCQA/astroid/blob/main/CONTRIBUTORS.txt
+
+"""
+Astroid hooks for numpy.core.einsumfunc module:
+https://github.com/numpy/numpy/blob/main/numpy/core/einsumfunc.py
+"""
+
+from astroid import nodes
+from astroid.brain.helpers import register_module_extender
+from astroid.builder import parse
+from astroid.manager import AstroidManager
+
+
+def numpy_core_einsumfunc_transform() -> nodes.Module:
+    return parse(
+        """
+    def einsum(*operands, out=None, optimize=False, **kwargs):
+        return numpy.ndarray([0, 0])
+    """
+    )
+
+
+register_module_extender(
+    AstroidManager(), "numpy.core.einsumfunc", numpy_core_einsumfunc_transform
+)

tests/unittest_brain_numpy_core_einsumfunc.py

@@ -0,0 +1,58 @@
+# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
+# For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
+# Copyright (c) https://github.com/PyCQA/astroid/blob/main/CONTRIBUTORS.txt
+
+from __future__ import annotations
+
+import pytest
+
+from astroid import builder, nodes
+
+try:
+    import numpy  # pylint: disable=unused-import
+
+    HAS_NUMPY = True
+except ImportError:
+    HAS_NUMPY = False
+
+
+def _inferred_numpy_func_call(func_name: str, *func_args: str) -> nodes.FunctionDef:
+    node = builder.extract_node(
+        f"""
+    import numpy as np
+    func = np.{func_name:s}
+    func({','.join(func_args):s})
+    """
+    )
+    return node.infer()
+
+
+@pytest.mark.skipif(not HAS_NUMPY, reason="This test requires the numpy library.")
+def test_numpy_function_calls_inferred_as_ndarray() -> None:
+    """
+    Test that calls to numpy functions are inferred as numpy.ndarray
+    """
+    method = "einsum"
+    inferred_values = list(
+        _inferred_numpy_func_call(method, "ii, np.arange(25).reshape(5, 5)")
+    )
+
+    assert len(inferred_values) == 1, f"Too much inferred value for {method:s}"
+    assert inferred_values[-1].pytype() in (
+        ".ndarray",
+    ), f"Illicit type for {method:s} ({inferred_values[-1].pytype()})"
+
+
+@pytest.mark.skipif(not HAS_NUMPY, reason="This test requires the numpy library.")
+def test_function_parameters() -> None:
+    instance = builder.extract_node(
+        f"""
+    import numpy
+    numpy.einsum #@
+    """
+    )
+    actual_args = instance.inferred()[0].args
+
+    assert actual_args.vararg == "operands"
+    assert [arg.name for arg in actual_args.kwonlyargs] == ["out", "optimize"]
+    assert actual_args.kwarg == "kwargs"