diff --git a/formulae/expr.py b/formulae/expr.py index 92391e6..ddb99de 100644 --- a/formulae/expr.py +++ b/formulae/expr.py @@ -19,7 +19,7 @@ def __repr__(self): # pragma: no cover def __str__(self): # pragma: no cover right = " ".join(str(self.value).splitlines(True)) - return f"Assign(name={self.name}, value={right}\n)" + return f"Assign(name={self.name}, value={right})" def accept(self, visitor): return visitor.visitAssignExpr(self) diff --git a/formulae/parser.py b/formulae/parser.py index e7ac14d..d707f69 100644 --- a/formulae/parser.py +++ b/formulae/parser.py @@ -195,6 +195,8 @@ def primary(self): # pylint: disable=too-many-return-statements return Literal(token.literal, lexeme=token.lexeme) elif self.match("BQNAME"): return QuotedName(self.previous()) + elif self.match("PYTHON_LITERAL"): + return Literal(self.previous().literal) elif self.match("LEFT_PAREN"): expr = self.expression() self.consume("RIGHT_PAREN", "Expect ')' after expression.") diff --git a/formulae/scanner.py b/formulae/scanner.py index c550333..e3de458 100644 --- a/formulae/scanner.py +++ b/formulae/scanner.py @@ -188,11 +188,17 @@ def number(self): self.add_token("NUMBER", token) + # pylint: disable=eval-used def identifier(self): # 'mod.function' is also an identifier while self.peek().isalnum() or self.peek() in [".", "_"]: self.advance() - self.add_token("IDENTIFIER") + + token = self.code[self.start : self.current] + if token in ("True", "False", "None"): # These are actually literals, not variable names + self.add_token("PYTHON_LITERAL", eval(token)) # Pass literals, not strings + else: + self.add_token("IDENTIFIER") def char(self): while self.peek() not in ["'", '"'] and not self.at_end(): diff --git a/formulae/terms/call_resolver.py b/formulae/terms/call_resolver.py index 71ff69b..23cd4ad 100644 --- a/formulae/terms/call_resolver.py +++ b/formulae/terms/call_resolver.py @@ -93,8 +93,6 @@ class LazyVariable: The name of the variable it represents. """ - BUILTINS = {"True": True, "False": False, "None": None} - def __init__(self, name): self.name = name @@ -129,16 +127,13 @@ def eval(self, data_mask, env): result: The value represented by this name in either the data mask or the environment. """ - if self.name in self.BUILTINS: - result = self.BUILTINS[self.name] - else: + try: + result = data_mask[self.name] + except KeyError: try: - result = data_mask[self.name] - except KeyError: - try: - result = env.namespace[self.name] - except KeyError as e: - raise e + result = env.namespace[self.name] + except KeyError as e: + raise e return result diff --git a/tests/test_parser.py b/tests/test_parser.py index 0965c46..7e2a501 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -264,3 +264,20 @@ def test_unclosed_function_call(): data = pd.DataFrame({"y": [1, 2], "x": [1, 2]}) with pytest.raises(ParseError, match="after arguments"): design_matrices("y ~ f(x", data) + + +def test_parse_python_literals(): + result = Parser(Scanner("f(x, True, y=False, z=None)").scan(False)).parse() + + assert isinstance(result.args[0], Variable) + + assert isinstance(result.args[1], Literal) + assert result.args[1].value is True + + assert isinstance(result.args[2], Assign) + assert isinstance(result.args[2].value, Literal) + assert result.args[2].value.value is False + + assert isinstance(result.args[3], Assign) + assert isinstance(result.args[3].value, Literal) + assert result.args[3].value.value is None \ No newline at end of file