Merge pull request #883 from hippo91/bug_pylint_4001

Bug pylint 4001

ChangeLog

@@ -6,6 +6,9 @@ astroid's ChangeLog
 What's New in astroid 2.5.0?
 ============================
 Release Date: TBA
+* Adds a brain for type object so that it is possible to write `type[int]` in annotation.
+
+  Fixes PyCQA/pylint#4001
 
 * Add ``__class_getitem__`` method to ``subprocess.Popen`` brain under Python 3.9 so that it is seen as subscriptable by pylint.
 

astroid/brain/brain_type.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+"""
+Astroid hooks for type support.
+
+Starting from python3.9, type object behaves as it had __class_getitem__ method.
+However it was not possible to simply add this method inside type's body, otherwise
+all types would also have this method. In this case it would have been possible
+to write str[int].
+Guido Van Rossum proposed a hack to handle this in the interpreter:
+https://github.com/python/cpython/blob/master/Objects/abstract.c#L186-L189
+
+This brain follows the same logic. It is no wise to add permanently the __class_getitem__ method 
+to the type object. Instead we choose to add it only in the case of a subscript node
+which inside name node is type. 
+Doing this type[int] is allowed whereas str[int] is not.
+
+Thanks to Lukasz Langa for fruitful discussion.
+"""
+import sys
+
+from astroid import MANAGER, extract_node, inference_tip, nodes
+
+
+PY39 = sys.version_info >= (3, 9)
+
+
+def _looks_like_type_subscript(node):
+    """
+    Try to figure out if a Name node is used inside a type related subscript
+
+    :param node: node to check
+    :type node: astroid.node_classes.NodeNG
+    :return: true if the node is a Name node inside a type related subscript
+    :rtype: bool
+    """
+    if isinstance(node, nodes.Name) and isinstance(node.parent, nodes.Subscript):
+        return node.name == "type"
+    return False
+
+
+def infer_type_sub(node, context=None):
+    """
+    Infer a type[...] subscript
+
+    :param node: node to infer
+    :type node: astroid.node_classes.NodeNG
+    :param context: inference context
+    :type context: astroid.context.InferenceContext
+    :return: the inferred node
+    :rtype: nodes.NodeNG
+    """
+    class_src = """
+    class type:
+        def __class_getitem__(cls, key):
+            return cls
+     """
+    node = extract_node(class_src)
+    return node.infer(context=context)
+
+
+if PY39:
+    MANAGER.register_transform(
+        nodes.Name, inference_tip(infer_type_sub), _looks_like_type_subscript
+    )

tests/unittest_brain.py

@@ -937,6 +937,42 @@ def test_sys_streams(self):
             self.assertEqual(raw.name, "FileIO")
 
 
+@test_utils.require_version("3.9")
+class TypeBrain(unittest.TestCase):
+    def test_type_subscript(self):
+        """
+        Check that type object has the __class_getitem__ method
+        when it is used as a subscript
+        """
+        src = builder.extract_node(
+            """
+            a: type[int] = int
+            """
+        )
+        val_inf = src.annotation.value.inferred()[0]
+        self.assertIsInstance(val_inf, astroid.ClassDef)
+        self.assertEqual(val_inf.name, "type")
+        meth_inf = val_inf.getattr("__class_getitem__")[0]
+        self.assertIsInstance(meth_inf, astroid.FunctionDef)
+
+    def test_invalid_type_subscript(self):
+        """
+        Check that a type (str for example) that inherits
+        from type does not have __class_getitem__ method even
+        when it is used as a subscript
+        """
+        src = builder.extract_node(
+            """
+            a: str[int] = "abc"
+            """
+        )
+        val_inf = src.annotation.value.inferred()[0]
+        self.assertIsInstance(val_inf, astroid.ClassDef)
+        self.assertEqual(val_inf.name, "str")
+        with self.assertRaises(astroid.exceptions.AttributeInferenceError):
+            meth_inf = val_inf.getattr("__class_getitem__")[0]
+
+
 @test_utils.require_version("3.6")
 class TypingBrain(unittest.TestCase):
     def test_namedtuple_base(self):