[FIX] Handle new added classes

[mohammedahmed18]

fixes CF-834

PR Type

Enhancement, Tests

---

Description

• Extract import-insertion helper function

• Support inserting newly added classes

• Update replacer to add unique classes

• Add tests for new class handling

---

Diagram Walkthrough

flowchart LR
  extractor["code_extractor.py: find_insertion_index_after_imports()"] -- "shared helper" --> replacer["code_replacer.py: OptimFunctionReplacer"]
  collector["OptimFunctionCollector: track new classes"] -- "collect new classes" --> replacer
  replacer -- "insert unique classes at top-level" --> module["Updated Module Body"]
  tests["tests/test_code_replacement.py"] -- "validate class + import insertion" --> module

Loading

File Walkthrough

	Relevant files
Enhancement	code_extractor.py Extract and reuse import insertion index finder                    codeflash/code_utils/code_extractor.py Add find_insertion_index_after_imports helper. Replace in-class private method with new helper usage. Preserve import scanning logic and early break on defs. +28/-27  code_replacer.py Support insertion of newly added classes                                  codeflash/code_utils/code_replacer.py Import shared insertion helper. Collect and store newly added classes. Insert unique classes after imports or last class. Adjust function insertion around classes/functions. +37/-8
Tests	test_code_replacement.py Add tests for new class insertion and expectations              tests/test_code_replacement.py Fix string literal formatting in tests. Update expected outputs to include new classes. Add new test for helper class extraction. +133/-10

---

[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 Index When inserting new classes, using max_class_index or find_insertion_index_after_imports may place classes after the last existing class, potentially interleaving with module-level code or docstrings. Validate that insertion before/after is correct for different module layouts (e.g., module docstring, variables, top-level statements) and that spacing/blank lines remain consistent. if self.new_classes: existing_class_names = {_node.name.value for _node in node.body if isinstance(_node, cst.ClassDef)} unique_classes = [ new_class for new_class in self.new_classes if new_class.name.value not in existing_class_names ] if unique_classes: new_classes_insertion_idx = max_class_index or find_insertion_index_after_imports(node) new_body = list( chain(node.body[:new_classes_insertion_idx], unique_classes, node.body[new_classes_insertion_idx:]) ) node = node.with_changes(body=new_body) if max_function_index is not None: Uniqueness Check Unique class detection uses only the class name; same-named classes from different scopes/files could be incorrectly deduplicated. Confirm deduping only targets true duplicates in the same module and does not drop intentional shadowed or redefined classes. if self.new_classes: existing_class_names = {_node.name.value for _node in node.body if isinstance(_node, cst.ClassDef)} unique_classes = [ new_class for new_class in self.new_classes if new_class.name.value not in existing_class_names ] if unique_classes: new_classes_insertion_idx = max_class_index or find_insertion_index_after_imports(node) Conditional Imports Detection of conditional imports assumes all bodies are SimpleStatementLine with only Import/ImportFrom. This may miss else/elif branches or nested structures. Verify it handles common patterns (if TYPE_CHECKING, try/except imports) without misplacing insertion index. 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 # Stop scanning once we reach a class or function definition. # Imports are supposed to be at the top of the file, but they can technically appear anywhere, even at the bottom of the file. # Without this check, a stray import later in the file # would incorrectly shift our insertion index below actual code definitions. if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)): break return insert_index

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Fix index falsy handling Using "or" here can yield 0 when there are no classes, which is correct, but it also treats index 0 as falsy and may incorrectly fall back when the last class is at index 0. Explicitly check for None to avoid misplacement of newly inserted classes. codeflash/code_utils/code_replacer.py [365] -new_classes_insertion_idx = max_class_index or find_insertion_index_after_imports(node) +new_classes_insertion_idx = max_class_index if max_class_index is not None else find_insertion_index_after_imports(node) Suggestion importance[1-10]: 9 __ Why: The current use of or will wrongly treat index 0 as falsy and fall back to find_insertion_index_after_imports, misplacing insertions when the last class is at top of file. Explicit None-check correctly preserves index 0 and improves correctness.	High
General	Robust conditional import detection This only detects unconditional import-only bodies; it misses imports nested under "else"/"elif" and will raise if "stmt.body" isn't an "IfBody" with ".body". Safely consider both "body" and "orelse" suites and guard attribute access. codeflash/code_utils/code_extractor.py [83-87] -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 = False +if isinstance(stmt, cst.If): + suites = [] + if hasattr(stmt, "body") and hasattr(stmt.body, "body"): + suites.append(stmt.body.body) + if hasattr(stmt, "orelse") and stmt.orelse and hasattr(stmt.orelse, "body") and hasattr(stmt.orelse.body, "body"): + suites.append(stmt.orelse.body.body) + if suites: + is_conditional_import = all( + isinstance(inner, cst.SimpleStatementLine) + and all(isinstance(child, (cst.Import, cst.ImportFrom)) for child in inner.body) + for suite in suites + for inner in suite + ) Suggestion importance[1-10]: 8 __ Why: The improved logic safely handles if bodies and optional orelse, avoiding attribute errors and recognizing imports under both branches; this aligns with the function’s goal and enhances robustness without altering intent.	Medium
	Improve class uniqueness check Name-only uniqueness can collide when a class with the same name exists but differs in content, silently dropping intended insertions. At minimum, also compare bases to reduce false positives when deciding uniqueness. codeflash/code_utils/code_replacer.py [359-363] -existing_class_names = {_node.name.value for _node in node.body if isinstance(_node, cst.ClassDef)} +def _class_sig(cls: cst.ClassDef) -> tuple[str, tuple[str, ...]]: + bases = tuple( + b.value.value if isinstance(b.value, cst.Name) else getattr(getattr(b.value, "attr", None), "value", "") + for b in cls.bases or [] + ) + return (cls.name.value, bases) -unique_classes = [ - new_class for new_class in self.new_classes if new_class.name.value not in existing_class_names -] +existing_class_sigs = {_class_sig(_node) for _node in node.body if isinstance(_node, cst.ClassDef)} +unique_classes = [cls for cls in self.new_classes if _class_sig(cls) not in existing_class_sigs] + Suggestion importance[1-10]: 6 __ Why: Expanding uniqueness to include base class signatures reduces false positives where same-named but different-inheritance classes exist; it’s reasonable and accurate, though still a heuristic and less critical than a functional bug.	Low