From 953477cfd7bc3b717d1117727dd26613bc35e7a4 Mon Sep 17 00:00:00 2001
From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com>
Date: Wed, 21 Feb 2024 17:33:08 +0000
Subject: [PATCH] fixup root names

---
 narwhals/pandas_like/utils.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/narwhals/pandas_like/utils.py b/narwhals/pandas_like/utils.py
index c491b6f79..1f111fc49 100644
--- a/narwhals/pandas_like/utils.py
+++ b/narwhals/pandas_like/utils.py
@@ -181,6 +181,7 @@ def evaluate_into_exprs(
 
 
 def register_expression_call(expr: ExprT, attr: str, *args: Any, **kwargs: Any) -> ExprT:
+    from narwhals.pandas_like.expr import Expr
     from narwhals.pandas_like.series import Series
 
     plx = get_namespace(expr)
@@ -201,11 +202,19 @@ def func(df: DataFrame | LazyFrame) -> list[Series]:
                 out.append(plx._create_series_from_scalar(_out, column))
         return out
 
+    root_names = expr._root_names
+    for arg in args:
+        if isinstance(arg, Expr):
+            root_names.extend(arg._root_names)
+    for arg in kwargs.values():
+        if isinstance(arg, Expr):
+            root_names.extend(arg._root_names)
+
     return plx._create_expr_from_callable(  # type: ignore[return-value]
         func,
         depth=expr._depth + 1,
         function_name=f"{expr._function_name}->{attr}",
-        root_names=expr._root_names,
+        root_names=root_names,
         output_names=expr._output_names,
     )