[CG-12126] Conditionals fix

codegen-examples/examples/snapshot_event_handler/pr_tasks.py

@@ -1,8 +1,6 @@
 import logging
-from codegen.agents.code_agent import CodeAgent
 from codegen.extensions.github.types.pull_request import PullRequestLabeledEvent
 
-from codegen.extensions.langchain.tools import GithubCreatePRCommentTool, GithubCreatePRReviewCommentTool, GithubViewPRTool
 from codegen.sdk.core.codebase import Codebase
 
 logging.basicConfig(level=logging.INFO, force=True)
@@ -11,25 +9,22 @@
 
 def lint_for_dev_import_violations(codebase: Codebase, event: PullRequestLabeledEvent):
     # Next.js codemod to detect imports of the react-dev-overlay module in production code
-    
+
     patch, commit_shas, modified_symbols = codebase.get_modified_symbols_in_pr(event.pull_request.number)
     modified_files = set(commit_shas.keys())
     from codegen.sdk.core.statements.if_block_statement import IfBlockStatement
 
-    DIR_NAME = 'packages/next/src/client/components/react-dev-overlay'
+    DIR_NAME = "packages/next/src/client/components/react-dev-overlay"
     directory = codebase.get_directory(DIR_NAME)
 
     violations = []
 
-
     false_operators = ["!=", "!=="]
     true_operators = ["===", "=="]
 
-
-
     def is_valid_block_expression(if_block: IfBlockStatement) -> bool:
         """Check if the if block has a valid environment check condition.
-        
+
         Valid conditions are:
         - process.env.NODE_ENV !== 'production'
         - process.env.NODE_ENV != 'production'
@@ -38,60 +33,57 @@ def is_valid_block_expression(if_block: IfBlockStatement) -> bool:
         """
         if not if_block.is_if_statement:
             return False
-            
+
         condition = if_block.condition
         # Get the operator without any whitespace
         operator = condition.operator[-1].source
-        
+
         # Check for non-production conditions
         if operator in false_operators and condition.source == f"process.env.NODE_ENV {operator} 'production'":
             return True
-            
+
         # Check for explicit development conditions
         if operator in true_operators and condition.source == f"process.env.NODE_ENV {operator} 'development'":
             return True
-
-        return False
 
+        return False
 
     def process_else_block_expression(else_block: IfBlockStatement) -> bool:
         """Check if the else block is valid by checking its parent if block.
-        
+
         Valid when the parent if block checks for production environment:
         - if (process.env.NODE_ENV === 'production') { ... } else { <our import> }
         - if (process.env.NODE_ENV == 'production') { ... } else { <our import> }
         """
         if not else_block.is_else_statement:
             return False
-            
+
         main_if = else_block._main_if_block
         if not main_if or not main_if.condition:
             return False
-            
+
         condition = main_if.condition
         operator = condition.operator[-1].source
-        
+
         # Valid if the main if block checks for production
         return operator in true_operators and condition.source == f"process.env.NODE_ENV {operator} 'production'"
 
-
     for file in directory.files(recursive=True):
         for imp in file.inbound_imports:
-
             if imp.file.filepath not in modified_files:
                 # skip if the import is not in the pull request's modified files
                 continue
             # Skip if the import is from within the target directory
             if directory.dirpath in imp.file.filepath:
                 # "✅ Valid import" if the import is within the target directory
                 continue
-                
+
             parent_if_block = imp.parent_of_type(IfBlockStatement)
-            
+
             # Check if import is in a valid environment check block
             if_block_valid = parent_if_block and is_valid_block_expression(parent_if_block)
             else_block_valid = parent_if_block and process_else_block_expression(parent_if_block)
-            
+
             # Skip if the import is properly guarded by environment checks
             if if_block_valid or else_block_valid:
                 # "✅ Valid import" these are guarded by non prod checks
@@ -102,7 +94,6 @@ def process_else_block_expression(else_block: IfBlockStatement) -> bool:
             violations.append(violation)
             logger.info(f"Found violation: {violation}")
 
-
     if violations:
         # Comment on PR with violations
         review_attention_message = "## Dev Import Violations Found\n\n"
@@ -111,4 +102,4 @@ def process_else_block_expression(else_block: IfBlockStatement) -> bool:
         review_attention_message += "\n\nPlease ensure that development imports are not imported in production code."
 
         # Create PR comment with the formatted message
-        codebase._op.create_pr_comment(event.pull_request.number, review_attention_message)
+        codebase._op.create_pr_comment(event.pull_request.number, review_attention_message)

src/codegen/sdk/core/expressions/name.py

@@ -50,6 +50,61 @@ def rename_if_matching(self, old: str, new: str):
         if self.source == old:
             self.edit(new)
 
+    @noapidoc
+    def _resolve_conditionals(self, conditional_parent: ConditionalBlock, name: str, original_resolved):
+        """Resolves name references within conditional blocks by traversing the conditional chain.
+
+        This method handles name resolution within conditional blocks (like if/elif/else statements) by:
+        1. Finding the appropriate search boundary based on the conditional block's position
+        2. Handling "fake" conditionals by traversing up the conditional chain
+        3. Yielding resolved names while respecting conditional block boundaries
+
+        Args:
+            conditional_parent (ConditionalBlock): The parent conditional block containing the name reference
+            name (str): The name being resolved
+            original_resolved: The originally resolved symbol that triggered this resolution
+
+        Yields:
+            Symbol | Import | WildcardImport: Resolved symbols found within the conditional blocks
+
+        Notes:
+            - A "fake" conditional is one where is_true_conditional() returns False
+            - The search_limit ensures we don't resolve names that appear after our target
+            - The method stops when it either:
+                a) Reaches the top of the conditional chain
+                b) Returns to the original conditional block
+                c) Can't find any more resolutions
+        """
+        search_limit = conditional_parent.start_byte_for_condition_block
+        if search_limit >= original_resolved.start_byte:
+            search_limit = original_resolved.start_byte - 1
+        if not conditional_parent.is_true_conditional(original_resolved):
+            # If it's a fake conditional we must skip any potential enveloping conditionals
+            def get_top_of_fake_chain(conditional, resolved, search_limit=0):
+                if skip_fake := conditional.parent_of_type(ConditionalBlock):
+                    if skip_fake.is_true_conditional(resolved):
+                        return skip_fake.start_byte_for_condition_block
+                    search_limit = skip_fake.start_byte_for_condition_block
+                    return get_top_of_fake_chain(skip_fake, conditional, search_limit)
+                return search_limit
+
+            if search_limit := get_top_of_fake_chain(conditional_parent, original_resolved):
+                search_limit = search_limit
+            else:
+                return
+
+        original_conditional = conditional_parent
+        while next_resolved := next(conditional_parent.resolve_name(name, start_byte=search_limit, strict=False), None):
+            yield next_resolved
+            next_conditional = next_resolved.parent_of_type(ConditionalBlock)
+            if not next_conditional or next_conditional == original_conditional:
+                return
+            search_limit = next_conditional.start_byte_for_condition_block
+            if next_conditional and not next_conditional.is_true_conditional(original_resolved):
+                pass
+            if search_limit >= next_resolved.start_byte:
+                search_limit = next_resolved.start_byte - 1
+
     @noapidoc
     @reader
     def resolve_name(self, name: str, start_byte: int | None = None, strict: bool = True) -> Generator["Symbol | Import | WildcardImport"]:
@@ -60,14 +115,8 @@ def resolve_name(self, name: str, start_byte: int | None = None, strict: bool =
             return
 
         if hasattr(resolved_name, "parent") and (conditional_parent := resolved_name.parent_of_type(ConditionalBlock)):
-            top_of_conditional = conditional_parent.start_byte
             if self.parent_of_type(ConditionalBlock) == conditional_parent:
                 # Use in the same block, should only depend on the inside of the block
                 return
-            for other_conditional in conditional_parent.other_possible_blocks:
-                if cond_name := next(other_conditional.resolve_name(name, start_byte=other_conditional.end_byte_for_condition_block), None):
-                    if cond_name.start_byte >= other_conditional.start_byte:
-                        yield cond_name
-                top_of_conditional = min(top_of_conditional, other_conditional.start_byte)
 
-            yield from self.resolve_name(name, top_of_conditional, strict=False)
+            yield from self._resolve_conditionals(conditional_parent=conditional_parent, name=name, original_resolved=resolved_name)

src/codegen/sdk/core/interfaces/conditional_block.py

@@ -6,7 +6,9 @@
 
 
 class ConditionalBlock(Statement, ABC):
-    """An interface for any code block that might not be executed in the code, e.g if block/else block/try block/catch block ect."""
+    """An interface for any code block that might not be executed in the code,
+    e.g if block/else block, try block/catch block ect.
+    """
 
     @property
     @abstractmethod
@@ -19,3 +21,17 @@ def other_possible_blocks(self) -> Sequence["ConditionalBlock"]:
     def end_byte_for_condition_block(self) -> int:
         """Returns the end byte for the specific condition block"""
         return self.end_byte
+
+    @property
+    @noapidoc
+    def start_byte_for_condition_block(self) -> int:
+        """Returns the start byte for the specific condition block"""
+        return self.start_byte
+
+    @noapidoc
+    def is_true_conditional(self, descendant) -> bool:
+        """Returns if this conditional is truly conditional,
+        this is necessary as an override for things like finally
+        statements that share a parent with try blocks
+        """
+        return True

src/codegen/sdk/core/interfaces/editable.py

@@ -1106,6 +1106,15 @@ def parent_of_types(self, types: set[type[T]]) -> T | None:
             return self.parent.parent_of_types(types)
         return None
 
+    def is_child_of(self, instance: Editable) -> bool:
+        """Checks if this node is a descendant of the given editable instance in the AST."""
+        if not self.parent:
+            return False
+        if self.parent is instance:
+            return True
+        else:
+            return self.parent.is_child_of(instance=instance)
+
     @reader
     def ancestors(self, type: type[T]) -> list[T]:
         """Find all ancestors of the node of the given type. Does not return itself"""

src/codegen/sdk/core/statements/if_block_statement.py

@@ -299,3 +299,10 @@ def end_byte_for_condition_block(self) -> int:
         if self.is_if_statement:
             return self.consequence_block.end_byte
         return self.end_byte
+
+    @property
+    @noapidoc
+    def start_byte_for_condition_block(self) -> int:
+        if self.is_if_statement:
+            return self.consequence_block.start_byte
+        return self.start_byte

src/codegen/sdk/core/statements/try_catch_statement.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar, override
 
 from codegen.sdk.core.interfaces.conditional_block import ConditionalBlock
 from codegen.sdk.core.interfaces.has_block import HasBlock
 from codegen.sdk.core.statements.block_statement import BlockStatement
 from codegen.sdk.core.statements.statement import StatementType
-from codegen.shared.decorators.docs import apidoc
+from codegen.shared.decorators.docs import apidoc, noapidoc
 
 if TYPE_CHECKING:
     from codegen.sdk.core.detached_symbols.code_block import CodeBlock
@@ -27,3 +27,26 @@ class TryCatchStatement(ConditionalBlock, BlockStatement[Parent], HasBlock, ABC,
 
     statement_type = StatementType.TRY_CATCH_STATEMENT
     finalizer: BlockStatement | None = None
+
+    @noapidoc
+    @override
+    def is_true_conditional(self, descendant) -> bool:
+        if descendant.is_child_of(self.finalizer):
+            return False
+        return True
+
+    @property
+    @noapidoc
+    def end_byte_for_condition_block(self) -> int:
+        if self.code_block:
+            return self.code_block.end_byte
+        else:
+            return self.end_byte
+
+    @property
+    @noapidoc
+    def start_byte_for_condition_block(self) -> int:
+        if self.code_block:
+            return self.code_block.start_byte - 1
+        else:
+            return self.start_byte

src/codegen/sdk/python/statements/try_catch_statement.py

@@ -104,11 +104,3 @@ def nested_code_blocks(self) -> list[PyCodeBlock]:
     @noapidoc
     def other_possible_blocks(self) -> Sequence[ConditionalBlock]:
         return self.except_clauses
-
-    @property
-    @noapidoc
-    def end_byte_for_condition_block(self) -> int:
-        if self.code_block:
-            return self.code_block.end_byte
-        else:
-            return self.end_byte

src/codegen/sdk/typescript/statements/try_catch_statement.py

@@ -36,7 +36,7 @@ class TSTryCatchStatement(TryCatchStatement["TSCodeBlock"], TSBlockStatement):
     def __init__(self, ts_node: TSNode, file_node_id: NodeId, ctx: CodebaseContext, parent: TSCodeBlock, pos: int | None = None) -> None:
         super().__init__(ts_node, file_node_id, ctx, parent, pos)
         if handler_node := self.ts_node.child_by_field_name("handler"):
-            self.catch = TSCatchStatement(handler_node, file_node_id, ctx, self.code_block)
+            self.catch = TSCatchStatement(handler_node, file_node_id, ctx, self)
         if finalizer_node := self.ts_node.child_by_field_name("finalizer"):
             self.finalizer = TSBlockStatement(finalizer_node, file_node_id, ctx, self.code_block)
 
@@ -102,11 +102,3 @@ def other_possible_blocks(self) -> Sequence[ConditionalBlock]:
             return [self.catch]
         else:
             return []
-
-    @property
-    @noapidoc
-    def end_byte_for_condition_block(self) -> int:
-        if self.code_block:
-            return self.code_block.end_byte
-        else:
-            return self.end_byte