CG-10731: Add ChainedAttribute.attribute_chain (#383)

- Adds ChainedAttribute.attribute_chain to be able view all members of a lengthy chained attribute in a list - Can easily now traverse a lengthy chained attribute - i.e calling attribute_chain on one of the functionCalls or one of the ChainedAttributes of `a().b().c.d.e()` -> `[FunctionCall(name=a), FunctionCall(name=b), Name(source=c), Name(source=d), FunctionCall(name=e)]` --------- Signed-off-by: dependabot[bot] <support@github.com>

Author: tawsifkamal

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

@@ -52,6 +52,13 @@ def __init__(self, node: TSNode, positional_idx: int, parent: FunctionCall) -> N
  self._name_node = self._parse_expression(name_node, default=Name)
  self._value_node = self._parse_expression(_value_node)
 
+ def __repr__(self) -> str:
+ keyword = f"keyword={self.name}, " if self.name else ""
+ value = f"value='{self.value}', " if self.value else ""
+ type = f"type={self.type}" if self.type else ""
+
+ return f"Argument({keyword}{value}{type})"
+
  @noapidoc
  @classmethod
  def from_argument_list(cls, node: TSNode, file_node_id: NodeId, G: CodebaseGraph, parent: FunctionCall) -> MultiExpression[Parent, Argument]:

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,33 @@ 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] | 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.
+
+ Returns:
+ FunctionCall[Parent] | 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
 
- # TODO: also define a successor?
+  return self.parent_of_type(FunctionCall)
 
  @property
  @noapidoc
@@ -581,6 +627,26 @@ def function_calls(self) -> list[FunctionCall]:
  # calls.append(call)
  return sort_editables(calls, dedupe=False)
 
+ @property
+ @reader
+ def attribute_chain(self) -> list[FunctionCall | Name]:
+ """Returns a list of elements in the chainedAttribute that the function call belongs in.
+
+ Breaks down chained expressions into individual components in order of appearance.
+ For example: `a.b.c().d` -> [Name("a"), Name("b"), FunctionCall("c"), Name("d")]
+
+ Returns:
+ list[FunctionCall | Name]: List of Name nodes (property access) and FunctionCall nodes (method calls)
+ """
+ if isinstance(self.get_name(), ChainedAttribute): # child is chainedAttribute. MEANING that this is likely in the middle or the last function call of a chained function call chain.
+ return self.get_name().attribute_chain
+ elif isinstance(
+ self.parent, ChainedAttribute
+ ): # does not have child chainedAttribute, but parent is chainedAttribute. MEANING that this is likely the TOP function call of a chained function call chain.
+ return self.parent.attribute_chain
+ else: # this is a standalone function call
+ return [self]
+
  @property
  @noapidoc
  def descendant_symbols(self) -> list[Importable]:
@@ -603,24 +669,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/expressions/chained_attribute.py

@@ -15,9 +15,11 @@
 from codegen.shared.decorators.docs import apidoc, noapidoc
 
 if TYPE_CHECKING:
+ from codegen.sdk.core.detached_symbols.function_call import FunctionCall
  from codegen.sdk.core.interfaces.has_name import HasName
  from codegen.sdk.core.interfaces.importable import Importable
 
+
 Object = TypeVar("Object", bound="Chainable")
 Attribute = TypeVar("Attribute", bound="Resolvable")
 Parent = TypeVar("Parent", bound="Expression")
@@ -74,6 +76,49 @@ def attribute(self) -> Attribute:
  """
  return self._attribute
 
+ @property
+ @reader
+ def attribute_chain(self) -> list["FunctionCall | Name"]:
+ """Returns a list of elements in a chained attribute expression.
+
+ Breaks down chained expressions into individual components in order of appearance.
+ For example: `a.b.c().d` -> [Name("a"), Name("b"), FunctionCall("c"), Name("d")]
+
+ Returns:
+ list[FunctionCall | Name]: List of Name nodes (property access) and FunctionCall nodes (method calls)
+ """
+ from codegen.sdk.core.detached_symbols.function_call import FunctionCall
+
+ ret = []
+ curr = self
+
+ # Traverse backwards in code (children of tree node)
+ while isinstance(curr, ChainedAttribute):
+ curr = curr.object
+
+ if isinstance(curr, FunctionCall):
+ ret.insert(0, curr)
+ curr = curr.get_name()
+ elif isinstance(curr, ChainedAttribute):
+ ret.insert(0, curr.attribute)
+
+ # This means that we have reached the base of the chain and the first item was an attribute (i.e a.b.c.func())
+ if isinstance(curr, Name) and not isinstance(curr.parent, FunctionCall):
+ ret.insert(0, curr)
+
+ curr = self
+
+ # Traversing forward in code (parents of tree node). Will add the current node as well
+ while isinstance(curr, ChainedAttribute) or isinstance(curr, FunctionCall):
+ if isinstance(curr, FunctionCall):
+ ret.append(curr)
+ elif isinstance(curr, ChainedAttribute) and not isinstance(curr.parent, FunctionCall):
+ ret.append(curr.attribute)
+
+ curr = curr.parent
+
+ return ret
+
  @property
  def object(self) -> Object:
  """Returns the object that contains the attribute being looked up.

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",

src/codegen/sdk/core/symbol_group.py

@@ -37,6 +37,9 @@ def __init__(self, file_node_id: NodeId, G: CodebaseGraph, parent: Parent, node:
  node = children[0].ts_node
  super().__init__(node, file_node_id, G, parent)
 
+ def __repr__(self) -> str:
+ return f"Collection({self.symbols})" if self.symbols is not None else super().__repr__()
+
  def _init_children(self): ...
 
  @repr_func # HACK

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 = """