From 68dda8aeb06f8e5b7316403f5c5224c26b77ff02 Mon Sep 17 00:00:00 2001 From: Simon Brugman Date: Sun, 10 Dec 2023 01:13:28 +0100 Subject: [PATCH] exclude other references --- refurb/checks/readability/fluid_interface.py | 16 ++++++++++++---- test/data/err_184.py | 5 +++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/refurb/checks/readability/fluid_interface.py b/refurb/checks/readability/fluid_interface.py index a1ac8a1..f0896e5 100644 --- a/refurb/checks/readability/fluid_interface.py +++ b/refurb/checks/readability/fluid_interface.py @@ -11,7 +11,7 @@ ReturnStmt, ) -from refurb.checks.common import check_block_like +from refurb.checks.common import check_block_like, ReadCountVisitor from refurb.error import Error from refurb.visitor import TraverserVisitor @@ -80,9 +80,17 @@ def check_call(node, name: str | None = None) -> bool: match node: # Single chain case CallExpr(callee=MemberExpr(expr=NameExpr(name=x), name=_)): - return name is None or name == x + if name is None or name == x: + # Exclude other references + x_expr = NameExpr(x) + x_expr.fullname = x + visitor = ReadCountVisitor(x_expr) + visitor.accept(node) + return visitor.read_count == 1 + return False + # Nested - case CallExpr(callee=MemberExpr(expr=call_node, name=_)): + case CallExpr(callee=MemberExpr(expr=call_node, name=y)): return check_call(call_node) return False @@ -92,7 +100,7 @@ class NameReferenceVisitor(TraverserVisitor): name: NameExpr referenced: bool - def __init__(self, name: NameExpr, stmt: Statement) -> None: + def __init__(self, name: NameExpr, stmt: Statement | None = None) -> None: super().__init__() self.name = name self.stmt = stmt diff --git a/test/data/err_184.py b/test/data/err_184.py index 43034cd..f9dc580 100644 --- a/test/data/err_184.py +++ b/test/data/err_184.py @@ -97,6 +97,11 @@ def assign_alternating(df, df2): # these will not +def _(x): + y = x.m() + return y.operation(*[v for v in y]) + + def assign_multiple_referenced(df, df2): df = df.select("column") result_df = df.select("another_column")