Skip to content

Commit

Permalink
Fixes regression in which builtins were detected as variables for expr
Browse files Browse the repository at this point in the history
transition

In the case of `cond.expr("len(foo) > d"), it would expect both len and
foo as keys to the condition. This ignores all builtins.
  • Loading branch information
elijahbenizzy committed May 22, 2024
1 parent 926b31e commit de323aa
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 3 deletions.
8 changes: 7 additions & 1 deletion burr/core/action.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import abc
import ast
import builtins
import copy
import inspect
import sys
Expand Down Expand Up @@ -204,14 +205,16 @@ def expr(expr: str) -> "Condition":
:return: A condition that evaluates the given expression
"""
tree = ast.parse(expr, mode="eval")
all_builtins = builtins.__dict__

# Visitor class to collect variable names
class NameVisitor(ast.NodeVisitor):
def __init__(self):
self.names = set()

def visit_Name(self, node):
self.names.add(node.id)
if node.id not in all_builtins:
self.names.add(node.id)

# Visit the nodes and collect variable names
visitor = NameVisitor()
Expand Down Expand Up @@ -261,6 +264,9 @@ def condition_func(state: State) -> bool:
name = f"{', '.join(f'{key}={value}' for key, value in sorted(kwargs.items()))}"
return Condition(keys, condition_func, name=name)

def __repr__(self):
return f"condition: {self._name}"

@classmethod
@property
def default(self) -> "Condition":
Expand Down
4 changes: 2 additions & 2 deletions tests/core/test_action.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,8 @@ def test_condition_expr():


def test_condition_expr_complex():
cond = Condition.expr("foo == 'bar' and baz == 'qux'")
assert cond.name == "foo == 'bar' and baz == 'qux'"
cond = Condition.expr("foo == 'bar' and len(baz) == 3")
assert cond.name == "foo == 'bar' and len(baz) == 3"
assert sorted(cond.reads) == ["baz", "foo"]
assert cond.run(State({"foo": "bar", "baz": "qux"})) == {Condition.KEY: True}
assert cond.run(State({"foo": "baz", "baz": "qux"})) == {Condition.KEY: False}
Expand Down

0 comments on commit de323aa

Please sign in to comment.