[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)

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

@@ -50,6 +50,62 @@ 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 +116,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

@@ -2,10 +2,13 @@
 from collections.abc import Sequence
 
 from codegen.sdk.core.statements.statement import Statement
+from codegen.shared.decorators.docs import noapidoc
 
 
 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
@@ -15,3 +18,17 @@ def other_possible_blocks(self) -> Sequence["ConditionalBlock"]:
     @property
     def end_byte_for_condition_block(self) -> int:
         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

@@ -297,3 +297,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/catch_statement.py

@@ -31,3 +31,4 @@ def __init__(self, ts_node: PyNode, file_node_id: NodeId, ctx: CodebaseContext,
     @property
     def other_possible_blocks(self) -> list[ConditionalBlock]:
         return [clause for clause in self.parent.except_clauses if clause != self] + [self.parent]
+

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

@@ -68,6 +68,7 @@ def _compute_dependencies(self, usage_type: UsageKind | None = None, dest: HasNa
         if self.finalizer:
             self.finalizer._compute_dependencies(usage_type, dest)
 
+
     @property
     @noapidoc
     def descendant_symbols(self) -> list[Importable]:
@@ -103,10 +104,3 @@ def nested_code_blocks(self) -> list[PyCodeBlock]:
     @property
     def other_possible_blocks(self) -> Sequence[ConditionalBlock]:
         return self.except_clauses
-
-    @property
-    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)
 
@@ -101,10 +101,3 @@ def other_possible_blocks(self) -> Sequence[ConditionalBlock]:
             return [self.catch]
         else:
             return []
-
-    @property
-    def end_byte_for_condition_block(self) -> int:
-        if self.code_block:
-            return self.code_block.end_byte
-        else:
-            return self.end_byte

tests/unit/codegen/sdk/python/statements/if_block_statement/test_if_block_statement_properties.py → tests/unit/codegen/sdk/python/statements/test_if_block_statement_properties.py

@@ -266,3 +266,56 @@ def test_if_else_reassigment_handling_nested_usage(tmpdir) -> None:
         second = file.symbols[1]
         assert len(first.usages) == 0
         assert second.usages[0].match == pyspark_arg
+
+
+def test_if_else_reassigment_inside_func_with_external_element(tmpdir) -> None:
+    content = """
+    PYSPARK="0"
+    def foo():
+        if True:
+            PYSPARK = True
+        else:
+            PYSPARK = False
+        print(PYSPARK)
+
+    """
+
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        funct_call = file.function_calls[0]
+        pyspark_arg = funct_call.args.children[0]
+        func = file.get_function("foo")
+        for assign in func.valid_symbol_names[:-1]:
+            assign.usages[0] == pyspark_arg
+
+
+
+def test_if_else_reassigment_handling_double_nested(tmpdir) -> None:
+    content = """
+        if False:
+            PYSPARK = "TEST1"
+        elif True:
+            PYSPARK = "TEST2"
+
+        if True:
+            PYSPARK = True
+        elif None:
+            if True:
+                PYSPARK = True
+            elif None:
+                if True:
+                    PYSPARK = True
+                elif None:
+                    PYSPARK = False
+
+        print(PYSPARK)
+    """
+
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        symbo = file.get_symbol("PYSPARK")
+        funct_call = file.function_calls[0]
+        pyspark_arg = funct_call.args.children[0]
+        for symb in file.symbols:
+            usage = symb.usages[0]
+            assert usage.match == pyspark_arg