adding attribute_chain

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

@@ -628,6 +628,18 @@ def function_calls(self) -> list[FunctionCall]:
         #     calls.append(call)
         return sort_editables(calls, dedupe=False)
 
+    @property
+    @reader
+    def attribute_chain(self) -> list[FunctionCall | Name]:
+        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]:

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,41 @@ def attribute(self) -> Attribute:
         """
         return self._attribute
 
+    @property
+    @reader
+    def attribute_chain(self) -> list["FunctionCall | Name"]:
+        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/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/expressions/test_chained_attribute_attribute_chain.py

@@ -0,0 +1,124 @@
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+
+
+def test_attribute_chain_query_builder(tmpdir) -> None:
+    # language=python
+    content = """
+def query():
+    # Test chained method calls with function at start
+    QueryBuilder().select("name", "age").from_table("users").where("age > 18").order_by("name")
+"""
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        query = file.get_function("query")
+        calls = query.function_calls
+        assert len(calls) == 5
+        order_by = calls[0]  # Last call in chain
+        where = calls[1]
+        from_table = calls[2]
+        select = calls[3]
+        query_builder = calls[4]  # First call in chain
+
+        # Test attribute chain from different positions
+        # From first call (QueryBuilder())
+        chain = query_builder.attribute_chain
+        assert len(chain) == 5
+        assert chain[0] == query_builder
+        assert chain[1] == select
+        assert chain[2] == from_table
+        assert chain[3] == where
+        assert chain[4] == order_by
+
+        # From middle call (from_table())
+        chain = from_table.attribute_chain
+        assert len(chain) == 5
+        assert chain[0] == query_builder
+        assert chain[1] == select
+        assert chain[2] == from_table
+        assert chain[3] == where
+        assert chain[4] == order_by
+
+        # From last call (order_by())
+        chain = order_by.attribute_chain
+        assert len(chain) == 5
+        assert chain[0] == query_builder
+        assert chain[1] == select
+        assert chain[2] == from_table
+        assert chain[3] == where
+        assert chain[4] == order_by
+
+
+def test_attribute_chain_mixed_properties(tmpdir) -> None:
+    # language=python
+    content = """
+def query():
+    # Test mix of properties and function calls
+    QueryBuilder().a.select("name", "age").from_table("users").where("age > 18").b.order_by("name").c
+"""
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        query = file.get_function("query")
+        calls = query.function_calls
+
+        # Get function calls in order
+        order_by = calls[0]  # Last function call
+        where = calls[1]
+        from_table = calls[2]
+        select = calls[3]
+        query_builder = calls[4]  # First function call
+
+        # Test from first call
+        chain = query_builder.attribute_chain
+        assert len(chain) == 8  # 5 function calls + 3 properties (a, b, c)
+        assert chain[0] == query_builder
+        assert chain[1].source == "a"  # Property
+        assert chain[2] == select
+        assert chain[3] == from_table
+        assert chain[4] == where
+        assert chain[5].source == "b"  # Property
+        assert chain[6] == order_by
+        assert chain[7].source == "c"  # Property
+
+
+def test_attribute_chain_only_properties(tmpdir) -> None:
+    # language=python
+    content = """
+def test():
+    # Test chain with only properties
+    a.b.c.func()
+"""
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        test = file.get_function("test")
+        calls = test.function_calls
+        assert len(calls) == 1
+        func = calls[0]
+
+        chain = func.attribute_chain
+        assert len(chain) == 4
+        assert chain[0].source == "a"
+        assert chain[1].source == "b"
+        assert chain[2].source == "c"
+        assert chain[3] == func
+
+
+def test_attribute_chain_nested_calls(tmpdir) -> None:
+    # language=python
+    content = """
+def test():
+    # Test nested function calls (not chained)
+    a(b(c()))
+"""
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}) as codebase:
+        file = codebase.get_file("test.py")
+        test = file.get_function("test")
+        calls = test.function_calls
+        assert len(calls) == 3
+        a = calls[0]
+        b = calls[1]
+        c = calls[2]
+
+        # Each call should have its own single-element chain
+        assert a.attribute_chain == [a]
+        assert b.attribute_chain == [b]
+        assert c.attribute_chain == [c]