Interpret True, False, and None as python literals (#107)

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)

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.")

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():

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
 
 

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