# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html # For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE # Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt from __future__ import annotations import unittest from typing import Any import astroid from astroid import MANAGER, builder, nodes from astroid.nodes.scoped_nodes import ClassDef try: import six # type: ignore[import] # pylint: disable=unused-import HAS_SIX = True except ImportError: HAS_SIX = False @unittest.skipUnless(HAS_SIX, "These tests require the six library") class SixBrainTest(unittest.TestCase): def test_attribute_access(self) -> None: ast_nodes = builder.extract_node( """ import six six.moves.http_client #@ six.moves.urllib_parse #@ six.moves.urllib_error #@ six.moves.urllib.request #@ """ ) assert isinstance(ast_nodes, list) http_client = next(ast_nodes[0].infer()) self.assertIsInstance(http_client, nodes.Module) self.assertEqual(http_client.name, "http.client") urllib_parse = next(ast_nodes[1].infer()) self.assertIsInstance(urllib_parse, nodes.Module) self.assertEqual(urllib_parse.name, "urllib.parse") urljoin = next(urllib_parse.igetattr("urljoin")) urlencode = next(urllib_parse.igetattr("urlencode")) self.assertIsInstance(urljoin, nodes.FunctionDef) self.assertEqual(urljoin.qname(), "urllib.parse.urljoin") self.assertIsInstance(urlencode, nodes.FunctionDef) self.assertEqual(urlencode.qname(), "urllib.parse.urlencode") urllib_error = next(ast_nodes[2].infer()) self.assertIsInstance(urllib_error, nodes.Module) self.assertEqual(urllib_error.name, "urllib.error") urlerror = next(urllib_error.igetattr("URLError")) self.assertIsInstance(urlerror, nodes.ClassDef) content_too_short = next(urllib_error.igetattr("ContentTooShortError")) self.assertIsInstance(content_too_short, nodes.ClassDef) urllib_request = next(ast_nodes[3].infer()) self.assertIsInstance(urllib_request, nodes.Module) self.assertEqual(urllib_request.name, "urllib.request") urlopen = next(urllib_request.igetattr("urlopen")) urlretrieve = next(urllib_request.igetattr("urlretrieve")) self.assertIsInstance(urlopen, nodes.FunctionDef) self.assertEqual(urlopen.qname(), "urllib.request.urlopen") self.assertIsInstance(urlretrieve, nodes.FunctionDef) self.assertEqual(urlretrieve.qname(), "urllib.request.urlretrieve") def test_from_imports(self) -> None: ast_node = builder.extract_node( """ from six.moves import http_client http_client.HTTPSConnection #@ """ ) inferred = next(ast_node.infer()) self.assertIsInstance(inferred, nodes.ClassDef) qname = "http.client.HTTPSConnection" self.assertEqual(inferred.qname(), qname) def test_from_submodule_imports(self) -> None: """Make sure ulrlib submodules can be imported from See pylint-dev/pylint#1640 for relevant issue """ ast_node = builder.extract_node( """ from six.moves.urllib.parse import urlparse urlparse #@ """ ) inferred = next(ast_node.infer()) self.assertIsInstance(inferred, nodes.FunctionDef) def test_with_metaclass_subclasses_inheritance(self) -> None: ast_node = builder.extract_node( """ class A(type): def test(cls): return cls class C: pass import six class B(six.with_metaclass(A, C)): pass B #@ """ ) inferred = next(ast_node.infer()) self.assertIsInstance(inferred, nodes.ClassDef) self.assertEqual(inferred.name, "B") self.assertIsInstance(inferred.bases[0], nodes.Call) ancestors = tuple(inferred.ancestors()) self.assertIsInstance(ancestors[0], nodes.ClassDef) self.assertEqual(ancestors[0].name, "C") self.assertIsInstance(ancestors[1], nodes.ClassDef) self.assertEqual(ancestors[1].name, "object") @staticmethod def test_six_with_metaclass_enum_ancestor() -> None: code = """ import six from enum import Enum, EnumMeta class FooMeta(EnumMeta): pass class Foo(six.with_metaclass(FooMeta, Enum)): #@ bar = 1 """ klass = astroid.extract_node(code) assert next(klass.ancestors()).name == "Enum" def test_six_with_metaclass_with_additional_transform(self) -> None: def transform_class(cls: Any) -> ClassDef: if cls.name == "A": cls._test_transform = 314 return cls MANAGER.register_transform(nodes.ClassDef, transform_class) try: ast_node = builder.extract_node( """ import six class A(six.with_metaclass(type, object)): pass A #@ """ ) inferred = next(ast_node.infer()) assert getattr(inferred, "_test_transform", None) == 314 finally: MANAGER.unregister_transform(nodes.ClassDef, transform_class)