Bug pylint 4326 (#1183)

* Adds unittest dealing with class subscript

* Adds support of type hints inside numpy's brains

* Adds unit test to check astroid does not crash if numpy is not available

Author: hippo91

ChangeLog

@@ -12,6 +12,10 @@ What's New in astroid 2.8.1?
 ============================
 Release date: TBA
 
+* Adds support of type hints inside numpy's brains.
+
+  Closes PyCQA/pylint#4326
+
 * Enable inference of dataclass import from pydantic.dataclasses.
   This allows the dataclasses brain to recognize pydantic dataclasses.
 

astroid/brain/brain_numpy_core_numerictypes.py

@@ -9,6 +9,7 @@
 # TODO(hippo91) : correct the methods signature.
 
 """Astroid hooks for numpy.core.numerictypes module."""
+from astroid.brain.brain_numpy_utils import numpy_supports_type_hints
 from astroid.brain.helpers import register_module_extender
 from astroid.builder import parse
 from astroid.manager import AstroidManager
@@ -19,9 +20,7 @@ def numpy_core_numerictypes_transform():
     #       According to numpy doc the generic object should expose
     #       the same API than ndarray. This has been done here partially
     #       through the astype method.
-    return parse(
-        """
-    # different types defined in numerictypes.py
+    generic_src = """
     class generic(object):
         def __init__(self, value):
             self.T = np.ndarray([0, 0])
@@ -106,8 +105,16 @@ def trace(self): return uninferable
         def transpose(self): return uninferable
         def var(self): return uninferable
         def view(self): return uninferable
-
-
+        """
+    if numpy_supports_type_hints():
+        generic_src += """
+        @classmethod
+        def __class_getitem__(cls, value):
+            return cls
+        """
+    return parse(
+        generic_src
+        + """
     class dtype(object):
         def __init__(self, obj, align=False, copy=False):
             self.alignment = None

astroid/brain/brain_numpy_ndarray.py

@@ -9,6 +9,7 @@
 
 
 """Astroid hooks for numpy ndarray class."""
+from astroid.brain.brain_numpy_utils import numpy_supports_type_hints
 from astroid.builder import extract_node
 from astroid.inference_tip import inference_tip
 from astroid.manager import AstroidManager
@@ -143,6 +144,12 @@ def transpose(self, *axes): return np.ndarray([0, 0])
         def var(self, axis=None, dtype=None, out=None, ddof=0, keepdims=False): return np.ndarray([0, 0])
         def view(self, dtype=None, type=None): return np.ndarray([0, 0])
     """
+    if numpy_supports_type_hints():
+        ndarray += """
+        @classmethod
+        def __class_getitem__(cls, value):
+            return cls
+        """
     node = extract_node(ndarray)
     return node.infer(context=context)
 

astroid/brain/brain_numpy_utils.py

@@ -8,9 +8,35 @@
 
 
 """Different utilities for the numpy brains"""
+from typing import Tuple
+
 from astroid.builder import extract_node
 from astroid.nodes.node_classes import Attribute, Import, Name, NodeNG
 
+# Class subscript is available in numpy starting with version 1.20.0
+NUMPY_VERSION_TYPE_HINTS_SUPPORT = ("1", "20", "0")
+
+
+def numpy_supports_type_hints() -> bool:
+    """
+    Returns True if numpy supports type hints
+    """
+    np_ver = _get_numpy_version()
+    return np_ver and np_ver > NUMPY_VERSION_TYPE_HINTS_SUPPORT
+
+
+def _get_numpy_version() -> Tuple[str, str, str]:
+    """
+    Return the numpy version number if numpy can be imported. Otherwise returns
+    ('0', '0', '0')
+    """
+    try:
+        import numpy  # pylint: disable=import-outside-toplevel
+
+        return tuple(numpy.version.version.split("."))
+    except ImportError:
+        return ("0", "0", "0")
+
 
 def infer_numpy_member(src, node, context=None):
     node = extract_node(src)

tests/unittest_brain_numpy_core_numerictypes.py

@@ -16,7 +16,12 @@
 except ImportError:
     HAS_NUMPY = False
 
-from astroid import builder, nodes
+from astroid import Uninferable, builder, nodes
+from astroid.brain.brain_numpy_utils import (
+    NUMPY_VERSION_TYPE_HINTS_SUPPORT,
+    _get_numpy_version,
+    numpy_supports_type_hints,
+)
 
 
 @unittest.skipUnless(HAS_NUMPY, "This test requires the numpy library.")
@@ -341,6 +346,88 @@ def test_datetime_astype_return(self):
             ),
         )
 
+    @unittest.skipUnless(
+        HAS_NUMPY and numpy_supports_type_hints(),
+        f"This test requires the numpy library with a version above {NUMPY_VERSION_TYPE_HINTS_SUPPORT}",
+    )
+    def test_generic_types_are_subscriptables(self):
+        """
+        Test that all types deriving from generic are subscriptables
+        """
+        for type_ in (
+            "bool_",
+            "bytes_",
+            "character",
+            "complex128",
+            "complex192",
+            "complex64",
+            "complexfloating",
+            "datetime64",
+            "flexible",
+            "float16",
+            "float32",
+            "float64",
+            "float96",
+            "floating",
+            "generic",
+            "inexact",
+            "int16",
+            "int32",
+            "int32",
+            "int64",
+            "int8",
+            "integer",
+            "number",
+            "signedinteger",
+            "str_",
+            "timedelta64",
+            "uint16",
+            "uint32",
+            "uint32",
+            "uint64",
+            "uint8",
+            "unsignedinteger",
+            "void",
+        ):
+            with self.subTest(type_=type_):
+                src = f"""
+                import numpy as np
+                np.{type_}[int]
+                """
+                node = builder.extract_node(src)
+                cls_node = node.inferred()[0]
+                self.assertIsInstance(cls_node, nodes.ClassDef)
+                self.assertEqual(cls_node.name, type_)
+
+
+@unittest.skipIf(
+    HAS_NUMPY, "Those tests check that astroid does not crash if numpy is not available"
+)
+class NumpyBrainUtilsTest(unittest.TestCase):
+    """
+    This class is dedicated to test that astroid does not crash
+    if numpy module is not available
+    """
+
+    def test_get_numpy_version_do_not_crash(self):
+        """
+        Test that the function _get_numpy_version doesn't crash even if numpy is not installed
+        """
+        self.assertEqual(_get_numpy_version(), ("0", "0", "0"))
+
+    def test_numpy_object_uninferable(self):
+        """
+        Test that in case numpy is not available, then a numpy object is uninferable
+        but the inference doesn't lead to a crash
+        """
+        src = """
+        import numpy as np
+        np.number[int]
+        """
+        node = builder.extract_node(src)
+        cls_node = node.inferred()[0]
+        self.assertIs(cls_node, Uninferable)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/unittest_brain_numpy_ndarray.py

@@ -16,7 +16,11 @@
 except ImportError:
     HAS_NUMPY = False
 
-from astroid import builder
+from astroid import builder, nodes
+from astroid.brain.brain_numpy_utils import (
+    NUMPY_VERSION_TYPE_HINTS_SUPPORT,
+    numpy_supports_type_hints,
+)
 
 
 @unittest.skipUnless(HAS_NUMPY, "This test requires the numpy library.")
@@ -161,6 +165,23 @@ def test_numpy_ndarray_attribute_inferred_as_ndarray(self):
                     msg=f"Illicit type for {attr_:s} ({inferred_values[-1].pytype()})",
                 )
 
+    @unittest.skipUnless(
+        HAS_NUMPY and numpy_supports_type_hints(),
+        f"This test requires the numpy library with a version above {NUMPY_VERSION_TYPE_HINTS_SUPPORT}",
+    )
+    def test_numpy_ndarray_class_support_type_indexing(self):
+        """
+        Test that numpy ndarray class can be subscripted (type hints)
+        """
+        src = """
+        import numpy as np
+        np.ndarray[int]
+        """
+        node = builder.extract_node(src)
+        cls_node = node.inferred()[0]
+        self.assertIsInstance(cls_node, nodes.ClassDef)
+        self.assertEqual(cls_node.name, "ndarray")
+
 
 if __name__ == "__main__":
     unittest.main()