Thanks to visit codestin.com
Credit goes to github.com

Skip to content

gh-123344: Add missing ast optimizations for PEP 696 #123377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Aug 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 95 additions & 9 deletions Lib/test/test_ast/test_ast.py
Original file line number Diff line number Diff line change
Expand Up @@ -3062,8 +3062,8 @@ class ASTOptimiziationTests(unittest.TestCase):
def wrap_expr(self, expr):
return ast.Module(body=[ast.Expr(value=expr)])

def wrap_for(self, for_statement):
return ast.Module(body=[for_statement])
def wrap_statement(self, statement):
return ast.Module(body=[statement])

def assert_ast(self, code, non_optimized_target, optimized_target):
non_optimized_tree = ast.parse(code, optimize=-1)
Expand All @@ -3090,16 +3090,16 @@ def assert_ast(self, code, non_optimized_target, optimized_target):
f"{ast.dump(optimized_tree)}",
)

def create_binop(self, operand, left=ast.Constant(1), right=ast.Constant(1)):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def create_binop(self, operand, left=ast.Constant(1), right=ast.Constant(1)):
def create_binop(self, operand, left=None, right=None):

please, don't use complex objects as defaults

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems fine as is; it's more convenient to write many of the remaining tests this way.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am in a purist camp :) But, this is minor, I agree. Since these objects aren't ever mutated, this does not introduce any problems.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I also don't like mutable defaults. But this one looks not very complex and it doesn't seem to be mutated anytime soon :) We can rewrite it, but the diff would be much larger

return ast.BinOp(left=left, op=self.binop[operand], right=right)

def test_folding_binop(self):
code = "1 %s 1"
operators = self.binop.keys()

def create_binop(operand, left=ast.Constant(1), right=ast.Constant(1)):
return ast.BinOp(left=left, op=self.binop[operand], right=right)

for op in operators:
result_code = code % op
non_optimized_target = self.wrap_expr(create_binop(op))
non_optimized_target = self.wrap_expr(self.create_binop(op))
optimized_target = self.wrap_expr(ast.Constant(value=eval(result_code)))

with self.subTest(
Expand All @@ -3111,7 +3111,7 @@ def create_binop(operand, left=ast.Constant(1), right=ast.Constant(1)):

# Multiplication of constant tuples must be folded
code = "(1,) * 3"
non_optimized_target = self.wrap_expr(create_binop("*", ast.Tuple(elts=[ast.Constant(value=1)]), ast.Constant(value=3)))
non_optimized_target = self.wrap_expr(self.create_binop("*", ast.Tuple(elts=[ast.Constant(value=1)]), ast.Constant(value=3)))
optimized_target = self.wrap_expr(ast.Constant(eval(code)))

self.assert_ast(code, non_optimized_target, optimized_target)
Expand Down Expand Up @@ -3222,12 +3222,12 @@ def test_folding_iter(self):
]

for left, right, ast_cls, optimized_iter in braces:
non_optimized_target = self.wrap_for(ast.For(
non_optimized_target = self.wrap_statement(ast.For(
target=ast.Name(id="_", ctx=ast.Store()),
iter=ast_cls(elts=[ast.Constant(1)]),
body=[ast.Pass()]
))
optimized_target = self.wrap_for(ast.For(
optimized_target = self.wrap_statement(ast.For(
target=ast.Name(id="_", ctx=ast.Store()),
iter=ast.Constant(value=optimized_iter),
body=[ast.Pass()]
Expand All @@ -3245,6 +3245,92 @@ def test_folding_subscript(self):

self.assert_ast(code, non_optimized_target, optimized_target)

def test_folding_type_param_in_function_def(self):
code = "def foo[%s = 1 + 1](): pass"

unoptimized_binop = self.create_binop("+")
unoptimized_type_params = [
("T", "T", ast.TypeVar),
("**P", "P", ast.ParamSpec),
("*Ts", "Ts", ast.TypeVarTuple),
]

for type, name, type_param in unoptimized_type_params:
result_code = code % type
optimized_target = self.wrap_statement(
ast.FunctionDef(
name='foo',
args=ast.arguments(),
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=ast.Constant(2))]
)
)
non_optimized_target = self.wrap_statement(
ast.FunctionDef(
name='foo',
args=ast.arguments(),
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=unoptimized_binop)]
)
)
self.assert_ast(result_code, non_optimized_target, optimized_target)

def test_folding_type_param_in_class_def(self):
code = "class foo[%s = 1 + 1]: pass"

unoptimized_binop = self.create_binop("+")
unoptimized_type_params = [
("T", "T", ast.TypeVar),
("**P", "P", ast.ParamSpec),
("*Ts", "Ts", ast.TypeVarTuple),
]

for type, name, type_param in unoptimized_type_params:
result_code = code % type
optimized_target = self.wrap_statement(
ast.ClassDef(
name='foo',
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=ast.Constant(2))]
)
)
non_optimized_target = self.wrap_statement(
ast.ClassDef(
name='foo',
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=unoptimized_binop)]
)
)
self.assert_ast(result_code, non_optimized_target, optimized_target)

def test_folding_type_param_in_type_alias(self):
code = "type foo[%s = 1 + 1] = 1"

unoptimized_binop = self.create_binop("+")
unoptimized_type_params = [
("T", "T", ast.TypeVar),
("**P", "P", ast.ParamSpec),
("*Ts", "Ts", ast.TypeVarTuple),
]

for type, name, type_param in unoptimized_type_params:
result_code = code % type
optimized_target = self.wrap_statement(
ast.TypeAlias(
name=ast.Name(id='foo', ctx=ast.Store()),
type_params=[type_param(name=name, default_value=ast.Constant(2))],
value=ast.Constant(value=1),
)
)
non_optimized_target = self.wrap_statement(
ast.TypeAlias(
name=ast.Name(id='foo', ctx=ast.Store()),
type_params=[type_param(name=name, default_value=unoptimized_binop)],
value=ast.Constant(value=1),
)
)
self.assert_ast(result_code, non_optimized_target, optimized_target)


if __name__ == "__main__":
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add AST optimizations for type parameter defaults.
3 changes: 3 additions & 0 deletions Python/ast_opt.c
Original file line number Diff line number Diff line change
Expand Up @@ -1087,10 +1087,13 @@ astfold_type_param(type_param_ty node_, PyArena *ctx_, _PyASTOptimizeState *stat
switch (node_->kind) {
case TypeVar_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.TypeVar.bound);
CALL_OPT(astfold_expr, expr_ty, node_->v.TypeVar.default_value);
break;
case ParamSpec_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.ParamSpec.default_value);
break;
case TypeVarTuple_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.TypeVarTuple.default_value);
break;
}
return 1;
Expand Down
Loading