forked from openkylin/astroid
252 lines
8.6 KiB
Python
252 lines
8.6 KiB
Python
# Copyright (c) 2015-2018, 2020 Claudiu Popa <pcmanticore@gmail.com>
|
|
# Copyright (c) 2015-2016 Ceridwen <ceridwenv@gmail.com>
|
|
# Copyright (c) 2016 Jakub Wilk <jwilk@jwilk.net>
|
|
# Copyright (c) 2018 Bryce Guinta <bryce.paul.guinta@gmail.com>
|
|
# Copyright (c) 2019 Ashley Whetter <ashley@awhetter.co.uk>
|
|
# Copyright (c) 2020-2021 hippo91 <guillaume.peillex@gmail.com>
|
|
# Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
|
|
# Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
|
|
|
|
# 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
|
|
|
|
|
|
import contextlib
|
|
import time
|
|
import unittest
|
|
from typing import Callable, Iterator, Optional
|
|
|
|
from astroid import MANAGER, builder, nodes, parse, transforms
|
|
from astroid.manager import AstroidManager
|
|
from astroid.nodes.node_classes import Call, Compare, Const, Name
|
|
from astroid.nodes.scoped_nodes import FunctionDef, Module
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def add_transform(
|
|
manager: AstroidManager,
|
|
node: type,
|
|
transform: Callable,
|
|
predicate: Optional[Callable] = None,
|
|
) -> Iterator:
|
|
manager.register_transform(node, transform, predicate)
|
|
try:
|
|
yield
|
|
finally:
|
|
manager.unregister_transform(node, transform, predicate)
|
|
|
|
|
|
class TestTransforms(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.transformer = transforms.TransformVisitor()
|
|
|
|
def parse_transform(self, code: str) -> Module:
|
|
module = parse(code, apply_transforms=False)
|
|
return self.transformer.visit(module)
|
|
|
|
def test_function_inlining_transform(self) -> None:
|
|
def transform_call(node: Call) -> Const:
|
|
# Let's do some function inlining
|
|
inferred = next(node.infer())
|
|
return inferred
|
|
|
|
self.transformer.register_transform(nodes.Call, transform_call)
|
|
|
|
module = self.parse_transform(
|
|
"""
|
|
def test(): return 42
|
|
test() #@
|
|
"""
|
|
)
|
|
|
|
self.assertIsInstance(module.body[1], nodes.Expr)
|
|
self.assertIsInstance(module.body[1].value, nodes.Const)
|
|
self.assertEqual(module.body[1].value.value, 42)
|
|
|
|
def test_recursive_transforms_into_astroid_fields(self) -> None:
|
|
# Test that the transformer walks properly the tree
|
|
# by going recursively into the _astroid_fields per each node.
|
|
def transform_compare(node: Compare) -> Const:
|
|
# Let's check the values of the ops
|
|
_, right = node.ops[0]
|
|
# Assume they are Consts and they were transformed before
|
|
# us.
|
|
return nodes.const_factory(node.left.value < right.value)
|
|
|
|
def transform_name(node: Name) -> Const:
|
|
# Should be Consts
|
|
return next(node.infer())
|
|
|
|
self.transformer.register_transform(nodes.Compare, transform_compare)
|
|
self.transformer.register_transform(nodes.Name, transform_name)
|
|
|
|
module = self.parse_transform(
|
|
"""
|
|
a = 42
|
|
b = 24
|
|
a < b
|
|
"""
|
|
)
|
|
|
|
self.assertIsInstance(module.body[2], nodes.Expr)
|
|
self.assertIsInstance(module.body[2].value, nodes.Const)
|
|
self.assertFalse(module.body[2].value.value)
|
|
|
|
def test_transform_patches_locals(self) -> None:
|
|
def transform_function(node: FunctionDef) -> None:
|
|
assign = nodes.Assign()
|
|
name = nodes.AssignName(name="value")
|
|
assign.targets = [name]
|
|
assign.value = nodes.const_factory(42)
|
|
node.body.append(assign)
|
|
|
|
self.transformer.register_transform(nodes.FunctionDef, transform_function)
|
|
|
|
module = self.parse_transform(
|
|
"""
|
|
def test():
|
|
pass
|
|
"""
|
|
)
|
|
|
|
func = module.body[0]
|
|
self.assertEqual(len(func.body), 2)
|
|
self.assertIsInstance(func.body[1], nodes.Assign)
|
|
self.assertEqual(func.body[1].as_string(), "value = 42")
|
|
|
|
def test_predicates(self) -> None:
|
|
def transform_call(node: Call) -> Const:
|
|
inferred = next(node.infer())
|
|
return inferred
|
|
|
|
def should_inline(node: Call) -> bool:
|
|
return node.func.name.startswith("inlineme")
|
|
|
|
self.transformer.register_transform(nodes.Call, transform_call, should_inline)
|
|
|
|
module = self.parse_transform(
|
|
"""
|
|
def inlineme_1():
|
|
return 24
|
|
def dont_inline_me():
|
|
return 42
|
|
def inlineme_2():
|
|
return 2
|
|
inlineme_1()
|
|
dont_inline_me()
|
|
inlineme_2()
|
|
"""
|
|
)
|
|
values = module.body[-3:]
|
|
self.assertIsInstance(values[0], nodes.Expr)
|
|
self.assertIsInstance(values[0].value, nodes.Const)
|
|
self.assertEqual(values[0].value.value, 24)
|
|
self.assertIsInstance(values[1], nodes.Expr)
|
|
self.assertIsInstance(values[1].value, nodes.Call)
|
|
self.assertIsInstance(values[2], nodes.Expr)
|
|
self.assertIsInstance(values[2].value, nodes.Const)
|
|
self.assertEqual(values[2].value.value, 2)
|
|
|
|
def test_transforms_are_separated(self) -> None:
|
|
# Test that the transforming is done at a separate
|
|
# step, which means that we are not doing inference
|
|
# on a partially constructed tree anymore, which was the
|
|
# source of crashes in the past when certain inference rules
|
|
# were used in a transform.
|
|
def transform_function(node: FunctionDef) -> Const:
|
|
if node.decorators:
|
|
for decorator in node.decorators.nodes:
|
|
inferred = next(decorator.infer())
|
|
if inferred.qname() == "abc.abstractmethod":
|
|
return next(node.infer_call_result())
|
|
return None
|
|
|
|
manager = MANAGER
|
|
with add_transform(manager, nodes.FunctionDef, transform_function):
|
|
module = builder.parse(
|
|
"""
|
|
import abc
|
|
from abc import abstractmethod
|
|
|
|
class A(object):
|
|
@abc.abstractmethod
|
|
def ala(self):
|
|
return 24
|
|
|
|
@abstractmethod
|
|
def bala(self):
|
|
return 42
|
|
"""
|
|
)
|
|
|
|
cls = module["A"]
|
|
ala = cls.body[0]
|
|
bala = cls.body[1]
|
|
self.assertIsInstance(ala, nodes.Const)
|
|
self.assertEqual(ala.value, 24)
|
|
self.assertIsInstance(bala, nodes.Const)
|
|
self.assertEqual(bala.value, 42)
|
|
|
|
def test_transforms_are_called_for_builtin_modules(self) -> None:
|
|
# Test that transforms are called for builtin modules.
|
|
def transform_function(node: FunctionDef) -> FunctionDef:
|
|
name = nodes.AssignName(name="value")
|
|
node.args.args = [name]
|
|
return node
|
|
|
|
manager = MANAGER
|
|
|
|
def predicate(node: FunctionDef) -> bool:
|
|
return node.root().name == "time"
|
|
|
|
with add_transform(manager, nodes.FunctionDef, transform_function, predicate):
|
|
builder_instance = builder.AstroidBuilder()
|
|
module = builder_instance.module_build(time)
|
|
|
|
asctime = module["asctime"]
|
|
self.assertEqual(len(asctime.args.args), 1)
|
|
self.assertIsInstance(asctime.args.args[0], nodes.AssignName)
|
|
self.assertEqual(asctime.args.args[0].name, "value")
|
|
|
|
def test_builder_apply_transforms(self) -> None:
|
|
def transform_function(node):
|
|
return nodes.const_factory(42)
|
|
|
|
manager = MANAGER
|
|
with add_transform(manager, nodes.FunctionDef, transform_function):
|
|
astroid_builder = builder.AstroidBuilder(apply_transforms=False)
|
|
module = astroid_builder.string_build("""def test(): pass""")
|
|
|
|
# The transform wasn't applied.
|
|
self.assertIsInstance(module.body[0], nodes.FunctionDef)
|
|
|
|
def test_transform_crashes_on_is_subtype_of(self) -> None:
|
|
# Test that we don't crash when having is_subtype_of
|
|
# in a transform, as per issue #188. This happened
|
|
# before, when the transforms weren't in their own step.
|
|
def transform_class(cls):
|
|
if cls.is_subtype_of("django.db.models.base.Model"):
|
|
return cls
|
|
return cls
|
|
|
|
self.transformer.register_transform(nodes.ClassDef, transform_class)
|
|
|
|
self.parse_transform(
|
|
"""
|
|
# Change environ to automatically call putenv() if it exists
|
|
import os
|
|
putenv = os.putenv
|
|
try:
|
|
# This will fail if there's no putenv
|
|
putenv
|
|
except NameError:
|
|
pass
|
|
else:
|
|
import UserDict
|
|
"""
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|