@@ -159,6 +159,8 @@ def _format_foldable_to_multiple_lines(
159
159
"annotation_args" : _format_args_to_multiple_lines ,
160
160
"inline_lambda" : _format_inline_lambda_to_multiple_lines ,
161
161
"lambda_header" : _format_lambda_header_to_multiple_lines ,
162
+ "inline_lambda_statements" : _format_inline_lambda_statements_to_multiple_lines ,
163
+ "return_stmt" : _format_return_stmt_to_multiple_lines ,
162
164
} # type: Dict[str, Callable]
163
165
return handlers [expression .data ](expression , expression_context , context )
164
166
@@ -570,11 +572,29 @@ def _format_inline_lambda_to_multiple_lines(
570
572
expression_context : ExpressionContext ,
571
573
context : Context ,
572
574
) -> FormattedLines :
575
+ expression_context_for_header = ExpressionContext (
576
+ expression_context .prefix_string , expression_context .prefix_line , "" , - 1
577
+ )
573
578
header_lines = _format_concrete_expression (
574
- inline_lambda .children [0 ], expression_context , context
579
+ inline_lambda .children [0 ], expression_context_for_header , context
580
+ )
581
+ _ , last_header_line = header_lines [- 1 ]
582
+ expression_context_for_statements = ExpressionContext (
583
+ f"{ last_header_line .strip ()} " ,
584
+ expression_context .prefix_line ,
585
+ expression_context .suffix_string ,
586
+ expression_context .suffix_line ,
587
+ )
588
+ fake_meta = Meta ()
589
+ fake_meta .line = inline_lambda .children [1 ].line
590
+ fake_meta .end_line = inline_lambda .children [- 1 ].end_line
591
+ fake_expression = Tree (
592
+ "inline_lambda_statements" , inline_lambda .children [1 :], fake_meta
593
+ )
594
+ statement_lines = _format_concrete_expression (
595
+ fake_expression , expression_context_for_statements , context
575
596
)
576
- # TODO: statements
577
- return header_lines [:- 1 ] + [(header_lines [- 1 ][0 ], header_lines [- 1 ][1 ] + " pass" )]
597
+ return header_lines [:- 1 ] + statement_lines
578
598
579
599
580
600
def _format_lambda_header_to_multiple_lines (
@@ -604,3 +624,53 @@ def _format_lambda_header_to_multiple_lines(
604
624
return _format_concrete_expression (
605
625
lambda_header .children [args_offset ], new_expression_context , context
606
626
)
627
+
628
+
629
+ def _format_inline_lambda_statements_to_multiple_lines (
630
+ inline_lambda_statements : Tree ,
631
+ expression_context : ExpressionContext ,
632
+ context : Context ,
633
+ ) -> FormattedLines :
634
+ lambda_statements = inline_lambda_statements .children
635
+ if len (lambda_statements ) == 1 :
636
+ return _format_concrete_expression (
637
+ lambda_statements [0 ], expression_context , context
638
+ )
639
+ expression_context_for_first_statement = ExpressionContext (
640
+ expression_context .prefix_string , expression_context .prefix_line , "" , - 1
641
+ )
642
+ first_statement_formatted_lines = _format_concrete_expression (
643
+ lambda_statements [0 ], expression_context_for_first_statement , context
644
+ )
645
+ _ , last_line = first_statement_formatted_lines [- 1 ]
646
+ remaining_statements_prefix = last_line .strip ()
647
+ remaining_statements_expression_context = ExpressionContext (
648
+ f"{ remaining_statements_prefix } ; " ,
649
+ expression_context .prefix_line ,
650
+ expression_context .suffix_string ,
651
+ expression_context .suffix_line ,
652
+ )
653
+ fake_meta = Meta ()
654
+ fake_meta .line = lambda_statements [1 ].line
655
+ fake_meta .end_line = lambda_statements [- 1 ].end_line
656
+ fake_expression = Tree ("inline_lambda_statements" , lambda_statements [1 :], fake_meta )
657
+ return first_statement_formatted_lines [:- 1 ] + _format_concrete_expression (
658
+ fake_expression , remaining_statements_expression_context , context
659
+ )
660
+
661
+
662
+ def _format_return_stmt_to_multiple_lines (
663
+ return_stmt : Tree ,
664
+ expression_context : ExpressionContext ,
665
+ context : Context ,
666
+ ) -> FormattedLines :
667
+ new_expression_context = ExpressionContext (
668
+ f"{ expression_context .prefix_string } return " ,
669
+ expression_context .prefix_line ,
670
+ expression_context .suffix_string ,
671
+ expression_context .suffix_line ,
672
+ )
673
+ formatted_lines , _ = format_expression (
674
+ return_stmt .children [0 ], new_expression_context , context
675
+ )
676
+ return formatted_lines
0 commit comments