lint

Author: tawsifkamal

src/codegen/sdk/core/detached_symbols/function_call.py

@@ -62,6 +62,28 @@ def __init__(self, node: TSNode, file_node_id: NodeId, G: CodebaseGraph, parent:
         args = [Argument(x, i, self) for i, x in enumerate(arg_list_node.named_children) if x.type != "comment"]
         self._arg_list = Collection(arg_list_node, self.file_node_id, self.G, self, children=args)
 
+    def __repr__(self) -> str:
+        """Custom string representation showing the function call chain structure.
+
+        Format: FunctionCall(name=current, pred=pred_name, succ=succ_name, base=base_name)
+
+        It will only print out predecessor, successor, and base that are of type FunctionCall. If it's a property, it will not be logged
+        """
+        # Helper to safely get name
+
+        # Get names for each part
+        parts = [f"name='{self.name}'"]
+
+        if self.predecessor and isinstance(self.predecessor, FunctionCall):
+            parts.append(f"predecessor=FunctionCall(name='{self.predecessor.name}')")
+
+        if self.successor and isinstance(self.successor, FunctionCall):
+            parts.append(f"successor=FunctionCall(name='{self.successor.name}')")
+
+        parts.append(f"filepath='{self.file.filepath}'")
+
+        return f"FunctionCall({', '.join(parts)})"
+
     @classmethod
     def from_usage(cls, node: Editable[Parent], parent: Parent | None = None) -> Self | None:
         """Creates a FunctionCall object from an Editable instance that represents a function call.
@@ -210,9 +232,36 @@ def predecessor(self) -> FunctionCall[Parent] | None:
             or if the predecessor is not a function call.
         """
         # Recursively travel down the tree to find the previous function call (child nodes are previous calls)
-        return self.call_chain[-2] if len(self.call_chain) > 1 else None
+        name = self.get_name()
+        while name:
+            if isinstance(name, FunctionCall):
+                return name
+            elif isinstance(name, ChainedAttribute):
+                name = name.object
+            else:
+                break
+        return None
+
+    @property
+    @reader
+    def successor(self) -> FunctionCall[Parent] | Name | None:
+        """Returns the next function call in a function call chain.
+
+        Returns the next function call in a function call chain. This method is useful for traversing function call chains
+        to analyze or modify sequences of chained function calls.
 
-    # TODO: also define a successor?
+        Args:
+            None
+
+        Returns:
+            FunctionCall[Parent] | Name | None: The next function call in the chain, or None if there is no successor
+            or if the successor is not a function call.
+        """
+        # this will avoid parent function calls in tree-sitter that are NOT part of the chained calls
+        if not isinstance(self.parent, ChainedAttribute):
+            return None
+
+        return self.parent_of_type(FunctionCall)
 
     @property
     @noapidoc
@@ -601,24 +650,35 @@ def register_api_call(self, url: str):
     @property
     @reader
     def call_chain(self) -> list[FunctionCall]:
-        """Returns a list of all function calls in this function call chain, including this call. Does not include calls made after this one."""
+        """Returns a list of all function  calls in this function call chain, including this call. Does not include calls made after this one."""
         ret = []
-        name = self.get_name()
-        while name:
-            if isinstance(name, FunctionCall):
-                ret.extend(name.call_chain)
-                break
-            elif isinstance(name, ChainedAttribute):
-                name = name.object
-            else:
-                break
+
+        # backward traversal
+        curr = self
+        pred = curr.predecessor
+        while pred is not None and isinstance(pred, FunctionCall):
+            ret.insert(0, pred)
+            pred = pred.predecessor
+
         ret.append(self)
+
+        # forward traversal
+        curr = self
+        succ = curr.successor
+        while succ is not None and isinstance(succ, FunctionCall):
+            ret.append(succ)
+            succ = succ.successor
+
         return ret
 
     @property
     @reader
     def base(self) -> Editable | None:
-        """Returns the base object of this function call chain."""
+        """Returns the base object of this function call chain.
+
+        Args:
+            Editable | None: The base object of this function call chain.
+        """
         name = self.get_name()
         while isinstance(name, ChainedAttribute):
             if isinstance(name.object, FunctionCall):

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

@@ -75,6 +75,10 @@ def _is_empty_container(text: str) -> bool:
     "resolved_types",
     "valid_symbol_names",
     "valid_import_names",
+    "predecessor",
+    "successor",
+    "base",
+    "call_chain",
     "code_block",
     "parent_statement",
     "symbol_usages",

tests/unit/codegen/sdk/python/detached_symbols/function_call/test_function_call.py

@@ -502,8 +502,8 @@ def baz():
 
         # Check call chain
         assert c.call_chain == [a, b, c]
-        assert b.call_chain == [a, b]
-        assert a.call_chain == [a]
+        assert b.call_chain == [a, b, c]
+        assert a.call_chain == [a, b, c]
 
         # Check base
         assert c.base == a.get_name()
@@ -530,15 +530,103 @@ def baz():
 
         # Check call chain
         assert c.call_chain == [a, b, c]
-        assert b.call_chain == [a, b]
-        assert a.call_chain == [a]
+        assert b.call_chain == [a, b, c]
+        assert a.call_chain == [a, b, c]
 
         # Check base
         assert c.base.source == "x"
         assert b.base.source == "x"
         assert a.base.source == "x"
 
 
+def test_function_call_chain_nested(tmpdir) -> None:
+    # language=python
+    content = """
+def foo():
+    # Nested function calls - each call should be independent
+    a(b(c()))
+"""
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        foo = file.get_function("foo")
+        calls = foo.function_calls
+        assert len(calls) == 3
+        a = calls[0]
+        b = calls[1]
+        c = calls[2]
+
+        # Each call should be independent - no predecessors
+        assert a.predecessor is None
+        assert b.predecessor is None
+        assert c.predecessor is None
+
+        # No successors since they're nested, not chained
+        assert a.successor is None
+        assert b.successor is None
+        assert c.successor is None
+
+        # Call chain for each should only include itself
+        assert a.call_chain == [a]
+        assert b.call_chain == [b]
+        assert c.call_chain == [c]
+
+        # Verify source strings are correct
+        assert a.source == "a(b(c()))"
+        assert b.source == "b(c())"
+        assert c.source == "c()"
+
+
+def test_function_call_chain_successor(tmpdir) -> None:
+    # language=python
+    content = """
+def foo():
+    a().b().c()
+
+def bat():
+    x.y.z.func()
+
+def baz():
+    x.a().y.b().z.c()
+"""
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+
+        # Check foo
+        foo = file.get_function("foo")
+        calls = foo.function_calls
+        assert len(calls) == 3
+        c = calls[0]
+        b = calls[1]
+        a = calls[2]
+
+        # Check successors
+        assert a.successor == b
+        assert b.successor == c
+        assert c.successor is None
+
+        # Check bat
+        bat = file.get_function("bat")
+        calls = bat.function_calls
+        assert len(calls) == 1
+        func = calls[0]
+
+        # No successor since it's a single function call
+        assert func.successor is None
+
+        # Check baz
+        baz = file.get_function("baz")
+        calls = baz.function_calls
+        assert len(calls) == 3
+        c = calls[0]
+        b = calls[1]
+        a = calls[2]
+
+        # Check successors
+        assert a.successor == b
+        assert b.successor == c
+        assert c.successor is None
+
+
 def test_function_call_chain_hard(tmpdir) -> None:
     # language=python
     content = """