diff options
| -rw-r--r-- | ChangeLog | 3 | ||||
| -rw-r--r-- | astroid/brain/brain_type.py | 64 | ||||
| -rw-r--r-- | tests/unittest_brain.py | 36 |
3 files changed, 103 insertions, 0 deletions
@@ -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. diff --git a/astroid/brain/brain_type.py b/astroid/brain/brain_type.py new file mode 100644 index 00000000..4e82813f --- /dev/null +++ b/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 + ) diff --git a/tests/unittest_brain.py b/tests/unittest_brain.py index 4717ab76..e6ff6928 100644 --- a/tests/unittest_brain.py +++ b/tests/unittest_brain.py @@ -937,6 +937,42 @@ class IOBrainTest(unittest.TestCase): 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): |