Skip to content

Commit

Permalink
Add formatting of simple lambdas
Browse files Browse the repository at this point in the history
  • Loading branch information
@Scony
Scony committed May 25, 2024
1 parent bb65adb commit 99aa761
Show file tree
Hide file tree
Showing 8 changed files with 140 additions and 162 deletions.
177 changes: 56 additions & 121 deletions gdtoolkit/formatter/expression.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import Dict, Callable, List, Optional
from importlib import import_module

from lark import Tree, Token
from lark.tree import Meta
Expand Down Expand Up @@ -149,6 +150,7 @@ def _format_foldable_to_multiple_lines(
"dict": _format_dict_to_multiple_lines,
"c_dict_element": _format_kv_pair_to_multiple_lines,
"eq_dict_element": _format_kv_pair_to_multiple_lines,
"lambda": _format_lambda_to_multiple_lines,
# fake expressions:
"func_args": _format_args_to_multiple_lines,
"func_arg_regular": _format_func_arg_to_multiple_lines,
Expand All @@ -165,33 +167,6 @@ def _format_foldable_to_multiple_lines(
),
"annotation": _format_annotation_to_multiple_lines,
"annotation_args": _format_args_to_multiple_lines,
"inline_lambda": _format_inline_lambda_to_multiple_lines,
"lambda_header": _format_lambda_header_to_multiple_lines,
"inline_lambda_statements": _format_inline_lambda_statements_to_multiple_lines,
"pass_stmt": _format_concrete_expression_to_single_line,
"return_stmt": lambda e, ec, c: _append_to_expression_context_and_pass_standalone(
"return ", e.children[0], ec, c
),
"expr_stmt": lambda e, ec, c: _format_standalone_expression(
e.children[0].children[0], ec, c
),
"func_var_stmt": lambda e, ec, c: _format_standalone_expression(
e.children[0], ec, c
),
"func_var_empty": _format_concrete_expression_to_single_line,
"func_var_assigned": lambda e, ec, c: _append_to_expression_context_and_pass_standalone(
f"var {expression_to_str(e.children[0])} = ", e.children[1], ec, c
),
"func_var_typed": _format_concrete_expression_to_single_line,
"func_var_typed_assgnd": lambda e, ec, c: _append_to_expression_context_and_pass_standalone(
f"var {expression_to_str(e.children[0])}: {expression_to_str(e.children[1])} = ",
e.children[2],
ec,
c,
),
"func_var_inf": lambda e, ec, c: _append_to_expression_context_and_pass_standalone(
f"var {expression_to_str(e.children[0])} := ", e.children[1], ec, c
),
"dot_chain": _format_dot_chain_to_multiple_lines,
"actual_getattr_call": _format_call_expression_to_multiple_lines,
"actual_subscr_expr": _format_subscription_to_multiple_lines,
Expand Down Expand Up @@ -645,100 +620,6 @@ def _format_annotation_to_multiple_lines(
)


def _format_inline_lambda_to_multiple_lines(
inline_lambda: Tree,
expression_context: ExpressionContext,
context: Context,
) -> FormattedLines:
expression_context_for_header = ExpressionContext(
expression_context.prefix_string, expression_context.prefix_line, "", -1
)
header_lines = _format_concrete_expression(
inline_lambda.children[0], expression_context_for_header, context
)
last_header_line_number, last_header_line = header_lines[-1]
assert last_header_line_number is not None
expression_context_for_statements = ExpressionContext(
f"{last_header_line.strip()} ",
last_header_line_number, # type:ignore
expression_context.suffix_string,
expression_context.suffix_line,
)
fake_meta = Meta()
fake_meta.line = get_line(inline_lambda.children[1])
fake_meta.end_line = get_end_line(inline_lambda.children[-1])
fake_expression = Tree(
"inline_lambda_statements", inline_lambda.children[1:], fake_meta
)
statement_lines = _format_concrete_expression(
fake_expression, expression_context_for_statements, context
)
return header_lines[:-1] + statement_lines


def _format_lambda_header_to_multiple_lines(
lambda_header: Tree,
expression_context: ExpressionContext,
context: Context,
) -> FormattedLines:
append_to_prefix = (
f"func {lambda_header.children[0].value}"
if isinstance(lambda_header.children[0], Token)
else "func"
)
args_offset = 1 if isinstance(lambda_header.children[0], Token) else 0
theres_something_after_args = len(lambda_header.children) > args_offset + 1
optional_type_hint = (
f" -> {lambda_header.children[args_offset+1]}"
if theres_something_after_args
else ""
)
prepend_to_suffix = f"{optional_type_hint}:"
new_expression_context = ExpressionContext(
f"{expression_context.prefix_string}{append_to_prefix}",
expression_context.prefix_line,
f"{prepend_to_suffix}{expression_context.suffix_string}",
expression_context.suffix_line,
)
return _format_concrete_expression(
lambda_header.children[args_offset], new_expression_context, context
)


def _format_inline_lambda_statements_to_multiple_lines(
inline_lambda_statements: Tree,
expression_context: ExpressionContext,
context: Context,
) -> FormattedLines:
lambda_statements = inline_lambda_statements.children
if len(lambda_statements) == 1:
return _format_concrete_expression(
lambda_statements[0], expression_context, context
)
expression_context_for_first_statement = ExpressionContext(
expression_context.prefix_string, expression_context.prefix_line, "", -1
)
first_statement_formatted_lines = _format_concrete_expression(
lambda_statements[0], expression_context_for_first_statement, context
)
last_line_number, last_line = first_statement_formatted_lines[-1]
assert last_line_number is not None
remaining_statements_prefix = last_line.strip()
remaining_statements_expression_context = ExpressionContext(
f"{remaining_statements_prefix} ; ",
last_line_number, # type: ignore
expression_context.suffix_string,
expression_context.suffix_line,
)
fake_meta = Meta()
fake_meta.line = get_line(lambda_statements[1])
fake_meta.end_line = get_end_line(lambda_statements[-1])
fake_expression = Tree("inline_lambda_statements", lambda_statements[1:], fake_meta)
return first_statement_formatted_lines[:-1] + _format_concrete_expression(
fake_expression, remaining_statements_expression_context, context
)


def _collapse_getattr_tree_to_dot_chain(expression: Tree) -> Tree:
reversed_dot_chain_children = [] # type: List[Node]
pending_getattr_call_to_match = None
Expand Down Expand Up @@ -883,3 +764,57 @@ def _format_dot_chain_to_multiple_lines_bottom_up(
fake_meta,
)
return _format_concrete_expression(new_actual_expr, expression_context, context)


def _format_lambda_to_multiple_lines(
a_lambda: Tree,
expression_context: ExpressionContext,
context: Context,
) -> FormattedLines:
assert expression_context.suffix_string == ""
expression_context_for_header = ExpressionContext(
expression_context.prefix_string, expression_context.prefix_line, "", -1
)
header_lines = _format_concrete_expression(
a_lambda.children[0], expression_context_for_header, context
)

block_module = import_module("gdtoolkit.formatter.block")
function_statement_module = import_module("gdtoolkit.formatter.function_statement")
child_context = context.create_child_context(expression_context.prefix_line)
(block_lines, _) = block_module.format_block(
a_lambda.children[1:],
function_statement_module.format_func_statement,
child_context,
)

return header_lines + block_lines


def _format_lambda_header_to_multiple_lines(
lambda_header: Tree,
expression_context: ExpressionContext,
context: Context,
) -> FormattedLines:
append_to_prefix = (
f"func {lambda_header.children[0].value}"
if isinstance(lambda_header.children[0], Token)
else "func"
)
args_offset = 1 if isinstance(lambda_header.children[0], Token) else 0
theres_something_after_args = len(lambda_header.children) > args_offset + 1
optional_type_hint = (
f" -> {lambda_header.children[args_offset+1]}"
if theres_something_after_args
else ""
)
prepend_to_suffix = f"{optional_type_hint}:"
new_expression_context = ExpressionContext(
f"{expression_context.prefix_string}{append_to_prefix}",
expression_context.prefix_line,
f"{prepend_to_suffix}{expression_context.suffix_string}",
expression_context.suffix_line,
)
return _format_concrete_expression(
lambda_header.children[args_offset], new_expression_context, context
)
34 changes: 7 additions & 27 deletions gdtoolkit/formatter/expression_to_str.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
remove_outer_parentheses,
has_trailing_comma,
)
from .function_statement_to_str import function_statement_to_str


def standalone_expression_to_str(expression: Node) -> str:
Expand Down Expand Up @@ -87,6 +88,8 @@ def expression_to_str(expression: Node) -> str:
[expression_to_str(n) for n in e.children]
),
"string_name": lambda e: f"&{expression_to_str(e.children[0])}",
"lambda": _lambda_to_str,
"lambda_header": _lambda_header_to_str,
# fake expressions:
"func_args": _args_to_str,
"func_arg_regular": lambda e: "{}{}".format(
Expand Down Expand Up @@ -121,30 +124,6 @@ def expression_to_str(expression: Node) -> str:
"trailing_comma": lambda _: "",
"annotation": _annotation_to_str,
"annotation_args": _annotation_args_to_str,
"inline_lambda": _inline_lambda_to_str,
"lambda_header": _lambda_header_to_str,
"inline_lambda_statements": lambda e: " ; ".join(
expression_to_str(statement) for statement in e.children
),
"pass_stmt": lambda _: "pass",
"return_stmt": lambda e: f"return {standalone_expression_to_str(e.children[0])}",
"expr_stmt": lambda e: f"{standalone_expression_to_str(e.children[0])}",
"func_var_stmt": lambda e: expression_to_str(e.children[0]),
"func_var_empty": lambda e: f"var {e.children[0].value}",
"func_var_assigned": lambda e: "var {} = {}".format(
e.children[0].value, standalone_expression_to_str(e.children[1])
),
"func_var_inf": lambda e: "var {} := {}".format(
e.children[0].value, standalone_expression_to_str(e.children[1])
),
"func_var_typed": lambda e: "var {}: {}".format(
e.children[0].value, standalone_expression_to_str(e.children[1])
),
"func_var_typed_assgnd": lambda e: "var {}: {} = {}".format(
e.children[0].value,
e.children[1].value,
standalone_expression_to_str(e.children[2]),
),
"non_foldable_dot_chain": lambda e: "".join(map(expression_to_str, e.children)),
"actual_getattr_call": _getattr_call_to_str,
"actual_subscr_expr": _subscription_to_str,
Expand Down Expand Up @@ -242,10 +221,11 @@ def _annotation_args_to_str(annotation: Tree) -> str:
return "({}{})".format(", ".join(elements), trailing_comma)


def _inline_lambda_to_str(inline_lambda: Tree) -> str:
fake_expression = Tree("inline_lambda_statements", inline_lambda.children[1:])
def _lambda_to_str(a_lambda: Tree) -> str:
assert len(a_lambda.children) == 2
return "{} {}".format(
expression_to_str(inline_lambda.children[0]), expression_to_str(fake_expression)
expression_to_str(a_lambda.children[0]),
function_statement_to_str(a_lambda.children[1]),
)


Expand Down
18 changes: 13 additions & 5 deletions gdtoolkit/formatter/expression_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,13 @@ def is_trailing_comma(expression: Node) -> bool:
def is_expression_forcing_multiple_lines(
expression: Node, standalone_comments: List[Optional[str]]
) -> bool:
if has_trailing_comma(expression):
return True
if _is_multiline_string(expression):
if has_trailing_comma(expression) or _is_multiline_string(expression):
return True
if isinstance(expression, Token):
return False
if _has_standalone_comments(expression, standalone_comments):
if _has_standalone_comments(
expression, standalone_comments
) or _is_multistatement_lambda(expression):
return True
for child in expression.children:
if is_expression_forcing_multiple_lines(child, standalone_comments):
Expand Down Expand Up @@ -82,10 +82,18 @@ def _is_multiline_string(expression: Node) -> bool:

def _has_standalone_comments(
expression: Tree, standalone_comments: List[Optional[str]]
):
) -> bool:
return any(
comment is not None
for comment in standalone_comments[
get_line(expression) : get_end_line(expression)
]
)


def _is_multistatement_lambda(expression: Tree) -> bool:
return (
isinstance(expression, Tree)
and expression.data == "lambda"
and len(expression.children) > 2
)
47 changes: 47 additions & 0 deletions gdtoolkit/formatter/function_statement_to_str.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
from importlib import import_module

from lark import Tree


def function_statement_to_str(statement: Tree) -> str:
expression_to_str_module = import_module("gdtoolkit.formatter.expression_to_str")
expression_to_str = expression_to_str_module.expression_to_str
standalone_expression_to_str = expression_to_str_module.standalone_expression_to_str
return {
"pass_stmt": lambda _: "pass",
"func_var_stmt": lambda s: function_statement_to_str(s.children[0]),
"const_stmt": _not_implemented,
"expr_stmt": lambda s: expression_to_str(s.children[0]), # TODO: standalone?
"return_stmt": _not_implemented,
"break_stmt": _not_implemented,
"breakpoint_stmt": _not_implemented,
"continue_stmt": _not_implemented,
"if_stmt": _not_implemented,
"while_stmt": _not_implemented,
"for_stmt": _not_implemented,
"for_stmt_typed": _not_implemented,
"match_stmt": _not_implemented,
"annotation": _not_implemented,
# statement fragments:
"func_var_empty": lambda s: f"var {s.children[0].value}",
"func_var_assigned": lambda s: "var {} = {}".format(
s.children[0].value, standalone_expression_to_str(s.children[1])
),
"func_var_inf": lambda s: "var {} := {}".format(
s.children[0].value, standalone_expression_to_str(s.children[1])
),
"func_var_typed": lambda s: "var {}: {}".format(
s.children[0].value, standalone_expression_to_str(s.children[1])
),
"func_var_typed_assgnd": lambda s: "var {}: {} = {}".format(
s.children[0].value,
s.children[1].value,
standalone_expression_to_str(s.children[2]),
),
"match_branch": _not_implemented,
"guarded_match_branch": _not_implemented,
}[statement.data](statement)


def _not_implemented(statement: Tree) -> str:
raise NotImplementedError
0 ...-output-pairs/simple_inline_lambdas.in.gd → ...r/input-output-pairs/simple_lambdas.in.gd
View file
File renamed without changes.
12 changes: 9 additions & 3 deletions ...output-pairs/simple_inline_lambdas.out.gd → .../input-output-pairs/simple_lambdas.out.gd
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,14 @@
func foo():
var x1 = func(): pass
var x2 = func(x): pass ; pass
var x3 = func(x: int): return 123 ; pass
var x4 = func bar(): pass ; 123
var x2 = func(x):
pass
pass
var x3 = func(x: int):
return 123
pass
var x4 = func bar():
pass
123
var x5 = func() -> int: 1
var x6 = func baz() -> int: pass
var x7 = func(): var x
Expand Down
Loading

0 comments on commit 99aa761

Please sign in to comment.