Skip to content

Commit

Permalink
tighten attribute detection
Browse files Browse the repository at this point in the history
  • Loading branch information
alexander-beedie committed Feb 16, 2024
1 parent 5868939 commit 6713cf6
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion py-polars/polars/utils/udfs.py
Original file line number Diff line number Diff line change
Expand Up @@ -659,6 +659,7 @@ def _matches(
*,
opnames: list[AbstractSet[str]],
argvals: list[AbstractSet[Any] | dict[Any, Any] | None] | None,
is_attr: bool = False,
) -> list[Instruction]:
"""
Check if a sequence of Instructions matches the specified ops/argvals.
Expand All @@ -671,9 +672,17 @@ def _matches(
The full opname sequence that defines a match.
argvals
Associated argvals that must also match (in same position as opnames).
is_attr
Indicate if the match is expected to represent attribute access.
"""
n_required_ops, argvals = len(opnames), argvals or []
instructions = self._instructions[idx : idx + n_required_ops]
idx_offset = idx + n_required_ops
if is_attr and (trailing_inst := self._instructions[idx_offset:1]):
# if there's a trailing CALL, it's an implicit method invocation
if trailing_inst.opname == "CALL":
return []

instructions = self._instructions[idx:idx_offset]
if len(instructions) == n_required_ops and all(
inst.opname in match_opnames
and (match_argval is None or inst.argval in match_argval)
Expand Down Expand Up @@ -717,6 +726,7 @@ def _rewrite_attrs(self, idx: int, updated_instructions: list[Instruction]) -> i
idx,
opnames=[{"LOAD_FAST"}, {"LOAD_ATTR"}],
argvals=[None, _PYTHON_ATTRS_MAP],
is_attr=True,
):
inst = matching_instructions[1]
expr_name = _PYTHON_ATTRS_MAP[inst.argval]
Expand Down

0 comments on commit 6713cf6

Please sign in to comment.