[Enhancement] Add global assignments after imports

[mohammedahmed18]

---

PR Type

Enhancement, Tests

---

Description

• Insert global assignments immediately after imports

• Add _find_insertion_index helper in code_extractor

• Update leave_Module to splice assignments via chain

• Revise tests to expect assignments after imports

---

Diagram Walkthrough

flowchart LR
  M["cst.Module body"] -- "find insertion index" --> I["_find_insertion_index"]
  I -- "insert assignments" --> U["updated module body"]

Loading

File Walkthrough

	Relevant files
Enhancement	code_extractor.py Insertion of assignments after imports                                      codeflash/code_utils/code_extractor.py Imported chain from itertools Added _find_insertion_index method Updated leave_Module to insert assignments after imports Revised blank-line logic after insertion +39/-12
Tests	test_code_replacement.py Adjust expected assignment placement                                          tests/test_code_replacement.py Added global assignment in expected code snippet Modified blank-line logic in test expectations +2/-2      test_multi_file_code_replacement.py Update multi-file test assignment placement                            tests/test_multi_file_code_replacement.py Added _TOKEN_SPLIT_RE assignment after imports Updated expected helper code with _translate_table +6/-8

---

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Insertion logic edge cases The insertion index loop breaks on any non-import statement, including blank lines or comments, which may lead to assignments being placed incorrectly. Consider skipping over EmptyLine and comment nodes before stopping. def _find_insertion_index(self, updated_node: cst.Module) -> int: insert_index = 0 for i, stmt in enumerate(updated_node.body): is_top_level_import = isinstance(stmt, cst.SimpleStatementLine) and any( isinstance(child, (cst.Import, cst.ImportFrom)) for child in stmt.body ) is_conditional_import = isinstance(stmt, cst.If) and all( isinstance(inner, cst.SimpleStatementLine) and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) for inner in stmt.body.body ) if is_top_level_import or is_conditional_import: insert_index = i + 1 else: # stop when we find the first non-import statement break return insert_index Conditional import detection The _find_insertion_index method only recognizes simple if blocks composed entirely of imports, but may miss elif or nested conditional patterns. Ensure all relevant import scenarios are captured. is_conditional_import = isinstance(stmt, cst.If) and all( isinstance(inner, cst.SimpleStatementLine) and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) for inner in stmt.body.body ) if is_top_level_import or is_conditional_import: Whitespace handling The logic that adds blank lines after assignments might introduce extra or missing spacing when the next statement already has leading empty lines. Validate and dedupe whitespace correctly. # Add a blank line after the last assignment if needed after_index = insert_index + len(assignment_lines) if after_index < len(new_statements): next_statement = new_statements[after_index] if not next_statement.leading_lines or not isinstance(next_statement.leading_lines[-1], cst.EmptyLine): new_statements[after_index] = next_statement.with_changes( leading_lines=[cst.EmptyLine(), *next_statement.leading_lines] )

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Skip docstrings and comments Add explicit handling for module docstrings and comments so they aren’t treated as non-import statements and keep insertions after them. Treat SimpleStatementLine with a SimpleString or comments as valid top-of-file lines. codeflash/code_utils/code_extractor.py [125-140] for i, stmt in enumerate(updated_node.body): - is_top_level_import = isinstance(stmt, cst.SimpleStatementLine) and any( - isinstance(child, (cst.Import, cst.ImportFrom)) for child in stmt.body + # skip module docstrings and comment lines + is_docstring = ( + isinstance(stmt, cst.SimpleStatementLine) + and isinstance(stmt.body[0], cst.Expr) + and isinstance(stmt.body[0].value, cst.SimpleString) ) - is_conditional_import = isinstance(stmt, cst.If) and all( - isinstance(inner, cst.SimpleStatementLine) - and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) - for inner in stmt.body.body + is_comment = ( + isinstance(stmt, cst.EmptyLine) + or ( + isinstance(stmt, cst.SimpleStatementLine) + and isinstance(stmt.body[0], cst.Comment) + ) + ) + if is_docstring or is_comment: + insert_index = i + 1 + continue + + is_top_level_import = ( + isinstance(stmt, cst.SimpleStatementLine) + and any(isinstance(child, (cst.Import, cst.ImportFrom)) for child in stmt.body) + ) + is_conditional_import = ( + isinstance(stmt, cst.If) + and all( + isinstance(inner, cst.SimpleStatementLine) + and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) + for inner in stmt.body.body + ) ) if is_top_level_import or is_conditional_import: insert_index = i + 1 else: break Suggestion importance[1-10]: 6 __ Why: The suggestion adds handling for module docstrings and comments to maintain correct insert position and is a useful improvement, though it isn’t critical.	Low
	Include else-branch imports Extend conditional import detection to cover else branches as well so that imports inside both if and else blocks are recognized. Include checks on stmt.orelse when it’s an Else suite. codeflash/code_utils/code_extractor.py [130-134] -is_conditional_import = isinstance(stmt, cst.If) and all( - isinstance(inner, cst.SimpleStatementLine) - and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) - for inner in stmt.body.body +is_conditional_import = ( + isinstance(stmt, cst.If) + and all( + isinstance(inner, cst.SimpleStatementLine) + and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) + for inner in stmt.body.body + ) + and ( + not stmt.orelse + or all( + isinstance(inner, cst.SimpleStatementLine) + and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) + for inner in stmt.orelse.body + ) + ) ) Suggestion importance[1-10]: 5 __ Why: Extending import detection to else branches handles more corner cases but has limited impact for most modules.	Low