[FIX] Respect top-level conditional imports e.g. (Conditions like TYPE_CHECKING & try/except &) (CF-382)

[mohammedahmed18]

User description

example: https://github.com/mohammedahmed18/pydantic-refi-exp/pull/33/files

PR Type

Bug fix, Tests

---

Description

• Introduce ConditionalImportCollector to track conditional imports

• Update add_needed_imports_from_module to skip those imports

• Refactor final module parsing to reuse parsed AST

• Add test_type_checking_imports for validation

---

Diagram Walkthrough

flowchart LR
  A["Parse destination module"] --> B["Collect conditional imports"]
  B --> C["Filter out conditional imports"]
  C --> D["Apply AddImportsVisitor and RemoveImportsVisitor"]

Loading

File Walkthrough

	Relevant files
Enhancement	code_extractor.py Skip conditional imports in extractor                                        codeflash/code_utils/code_extractor.py Added ConditionalImportCollector visitor class Prevent adding imports found in top-level conditionals Integrated collector into add_needed_imports_from_module Reuse parsed_dst_module for transformations +76/-6
Tests	test_code_replacement.py Add conditional import skipping test                                          tests/test_code_replacement.py Added import re for regex assertions Added test_type_checking_imports function Assert conditional imports are not re-added +137/-0

---

---

PR Type

Bug fix, Tests

---

Description

• Add ConditionalImportCollector for top-level conditionals

• Skip conditional imports in add_needed_imports_from_module

• Reuse parsed_dst_module for transformations

• Add test_type_checking_imports for validation

---

Diagram Walkthrough

flowchart LR
  ParseDST["Parse destination module AST"]
  CollectCond["ConditionalImportCollector collects imports"]
  Filter["Filter out conditional imports"]
  ApplyVisitors["Apply and remove imports"]

  ParseDST -- "visit" --> CollectCond
  CollectCond -- "collected imports" --> Filter
  Filter -- "skip imports" --> ApplyVisitors

Loading

File Walkthrough

	Relevant files
Bug fix	code_extractor.py Skip conditional top-level imports                                              codeflash/code_utils/code_extractor.py Added ConditionalImportCollector visitor class Collected imports inside top-level if/try blocks Integrated collector into add_needed_imports_from_module Reused parsed_dst_module for transformation +76/-6
Tests	test_code_replacement.py Add test for conditional imports skipping                                tests/test_code_replacement.py Imported re for regex validations Added test_type_checking_imports function Validates skipping of conditional imports +137/-0

---

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit f2e5cb2)

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 Import alias mismatch ConditionalImportCollector stores alias imports as "module.alias", but the filtering loops compare module names and alias names separately. This mismatch can cause conditional imports with aliases not to be skipped correctly. if isinstance(child, cst.Import): for alias in child.names: module = self.get_full_dotted_name(alias.name) asname = alias.asname.name.value if alias.asname else alias.name.value self.imports.add(module if module == asname else f"{module}.{asname}") Incomplete Try handler coverage _collect_imports_from_block only inspects the try block body, ignoring imports in except and else handlers. Conditional imports in those blocks won't be collected and may be reapplied. def _collect_imports_from_block(self, block: cst.IndentedBlock) -> None: for statement in block.body: if isinstance(statement, cst.SimpleStatementLine): for child in statement.body: if isinstance(child, cst.Import): for alias in child.names: module = self.get_full_dotted_name(alias.name) asname = alias.asname.name.value if alias.asname else alias.name.value self.imports.add(module if module == asname else f"{module}.{asname}") elif isinstance(child, cst.ImportFrom): if child.module is None: continue module = self.get_full_dotted_name(child.module) for alias in child.names: if isinstance(alias, cst.ImportAlias): name = alias.name.value asname = alias.asname.name.value if alias.asname else name self.imports.add(f"{module}.{asname}") def visit_Module(self, node: cst.Module) -> None: self.depth = 0 def visit_FunctionDef(self, node: cst.FunctionDef) -> None: self.depth += 1 def leave_FunctionDef(self, node: cst.FunctionDef) -> None: self.depth -= 1 def visit_ClassDef(self, node: cst.ClassDef) -> None: self.depth += 1 def leave_ClassDef(self, node: cst.ClassDef) -> None: self.depth -= 1 def visit_If(self, node: cst.If) -> None: if self.depth == 0: self._collect_imports_from_block(node.body) def visit_Try(self, node: cst.Try) -> None: if self.depth == 0: self._collect_imports_from_block(node.body) Faulty regex assertions The regex patterns use word boundaries after underscores (e.g., \b after aiohttp_), which will never match correctly, making these assertions ineffective. assert not re.search(r"^import requests\b", new_code, re.MULTILINE) # conditional simple import: import <name> assert not re.search(r"^import aiohttp as aiohttp_\b", new_code, re.MULTILINE) # conditional alias import: import <name> as <alias> assert not re.search(r"^from math import pi as PI, sin as sine\b", new_code, re.MULTILINE) # conditional multiple aliases imports

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to f2e5cb2
Explore these optional code suggestions:

Category	Suggestion	Impact
General	include else block imports Also collect imports from the if node’s orelse block so that conditional imports in else/elif branches are captured and skipped. codeflash/code_utils/code_extractor.py [247-249] def visit_If(self, node: cst.If) -> None: if self.depth == 0: self._collect_imports_from_block(node.body) + if node.orelse: + self._collect_imports_from_block(node.orelse) Suggestion importance[1-10]: 7 __ Why: The collector currently skips node.orelse imports, so adding them prevents missing conditional branches and ensures all conditional imports are skipped.	Medium
	include try/except/final imports Extend visit_Try to also collect imports from exception handlers, orelse, and finalbody blocks so all conditional imports within the try/except/finally are detected. codeflash/code_utils/code_extractor.py [251-253] def visit_Try(self, node: cst.Try) -> None: if self.depth == 0: self._collect_imports_from_block(node.body) + for handler in node.handlers: + self._collect_imports_from_block(handler.body) + if node.orelse: + self._collect_imports_from_block(node.orelse) + if node.finalbody: + self._collect_imports_from_block(node.finalbody) Suggestion importance[1-10]: 7 __ Why: Without collecting from handlers, orelse, and finalbody, some conditional imports inside exception blocks aren’t recognized, so extending coverage is valuable.	Medium
	guard empty module names Guard against empty module names returned by get_full_dotted_name to avoid adding malformed entries; skip the alias if module is empty. codeflash/code_utils/code_extractor.py [218-220] module = self.get_full_dotted_name(alias.name) +if not module: + continue asname = alias.asname.name.value if alias.asname else alias.name.value self.imports.add(module if module == asname else f"{module}.{asname}") Suggestion importance[1-10]: 5 __ Why: Skipping empty strings from get_full_dotted_name avoids malformed entries; this is a good safety check though unlikely often triggered.	Low

---

Previous suggestions

Suggestions up to commit 24f2785

Category	Suggestion	Impact
Possible issue	Collect imports in else blocks Extend visit_If to also collect imports in the else and elif branches so conditional imports there are respected. This prevents re-adding imports that only appear in an else block. codeflash/code_utils/code_extractor.py [247-249] def visit_If(self, node: cst.If) -> None: if self.depth == 0: self._collect_imports_from_block(node.body) + # also collect in else/elif blocks + if node.orelse: + orelse = node.orelse.body if isinstance(node.orelse, cst.Else) else node.orelse + self._collect_imports_from_block(orelse) Suggestion importance[1-10]: 5 __ Why: Extending visit_If to include else/elif branches helps avoid re-adding imports defined there and improves accuracy without conflicting existing logic.	Low
	Collect imports in except/finally Also traverse except handlers and the finally block in top-level try statements to capture conditional imports there. This ensures imports in those blocks aren’t reintroduced. codeflash/code_utils/code_extractor.py [251-253] def visit_Try(self, node: cst.Try) -> None: if self.depth == 0: + # try block self._collect_imports_from_block(node.body) + # except handlers + for handler in node.handlers: + self._collect_imports_from_block(handler.body) + # finally block + if node.finalbody: + self._collect_imports_from_block(node.finalbody) Suggestion importance[1-10]: 5 __ Why: Traversing except handlers and the finally block captures all conditional imports in top-level try statements, enhancing completeness.	Low
General	Catch all parse errors Broaden the exception catch to Exception when parsing and visiting the destination module so unexpected errors don’t bubble up. This keeps the function safe from all parsing failures. codeflash/code_utils/code_extractor.py [392-397] try: parsed_dst_module = cst.parse_module(dst_module_code) parsed_dst_module.visit(cond_import_collector) -except cst.ParserSyntaxError as e: - logger.exception(f"Syntax error in destination module code: {e}") - return dst_module_code # Return the original code if there's a syntax error +except Exception as e: + logger.exception(f"Error collecting conditional imports: {e}") + return dst_module_code Suggestion importance[1-10]: 4 __ Why: Broadening to Exception ensures unexpected parsing or visiting errors are handled, but risks hiding non-syntax issues and has limited impact.	Low

[github-actions]

Persistent review updated to latest commit f2e5cb2