(core) Ensure formulas return something and don't assign to attributes of rec

Summary: To help with mistakes in formulas, forbid assigning to attributes of `rec` (e.g. `$foo = 1` which should probably be `==`) and ensure that there is at least one `return` in the formula (after maybe adding an implicit one at the end).

Test Plan: Extended Python unit test, updated tests which were missing return.

Reviewers: georgegevoian

Reviewed By: georgegevoian

Subscribers: dsagal

Differential Revision: https://phab.getgrist.com/D3439

sandbox/grist/codebuilder.py

@@ -1,5 +1,6 @@
 import ast
 import contextlib
+import itertools
 import re
 import six
 
@@ -48,13 +49,6 @@ def make_formula_body(formula, default_value, assoc_value=None):
   except SyntaxError as e:
     return textbuilder.Text(_create_syntax_error_code(tmp_formula, formula, e))
 
-  # Parse formula and generate error code on assignment to rec
-  with use_inferences(InferRecAssignment):
-    try:
-      astroid.parse(tmp_formula.get_text())
-    except SyntaxError as e:
-      return textbuilder.Text(_create_syntax_error_code(tmp_formula, formula, e))
-
   # Once we have a tree, go through it and create a subset of the dollar patches that are actually
   # relevant. E.g. this is where we'll skip the "$foo" patches that appear in strings or comments.
   patches = []
@@ -86,6 +80,19 @@ def make_formula_body(formula, default_value, assoc_value=None):
   elif last_statement is None:
     # If we have an empty body (e.g. just a comment), add a 'pass' at the end.
     patches.append(textbuilder.make_patch(formula, len(formula), len(formula), '\npass'))
+  elif not any(
+      # Raise an error if the user forgot to return anything. For performance:
+      # - Use type() instead of isinstance()
+      # - Check last_statement first to try avoiding walking the tree
+      type(node) == ast.Return  # pylint: disable=unidiomatic-typecheck
+      for node in itertools.chain([last_statement], ast.walk(atok.tree))
+  ):
+    message = "No `return` statement, and the last line isn't an expression."
+    if isinstance(last_statement, ast.Assign):
+      message += " If you want to check for equality, use `==` instead of `=`."
+    error = SyntaxError(message,
+                        ('<string>', 1, 1, ""))
+    return textbuilder.Text(_create_syntax_error_code(tmp_formula, formula, error))
 
   # Apply the new set of patches to the original formula to get the real output.
   final_formula = textbuilder.Replacer(formula_builder_text, patches)
@@ -93,10 +100,13 @@ def make_formula_body(formula, default_value, assoc_value=None):
   # Try parsing again before returning it just in case we have new syntax errors. These are
   # possible in cases when a single token ('DOLLARfoo') is valid but an expression ('rec.foo') is
   # not, e.g. `foo($bar=1)` or `def $foo()`.
-  try:
-    atok = asttokens.ASTTokens(final_formula.get_text(), parse=True)
-  except SyntaxError as e:
-    return textbuilder.Text(_create_syntax_error_code(final_formula, formula, e))
+  # Also check for common mistakes: assigning to `rec` or its attributes (e.g. `$foo = 1`).
+  with use_inferences(InferRecAssignment, InferRecAttrAssignment):
+    try:
+      astroid.parse(final_formula.get_text())
+    except (astroid.AstroidSyntaxError, SyntaxError) as e:
+      error = getattr(e, "error", e)  # extract SyntaxError from AstroidSyntaxError
+      return textbuilder.Text(_create_syntax_error_code(final_formula, formula, error))
 
   # We return the text-builder object whose .get_text() is the final formula.
   return final_formula
@@ -268,6 +278,23 @@ class InferRecAssignment(InferenceTip):
   def infer(cls, node, context):
     raise NotImplementedError()
 
+class InferRecAttrAssignment(InferenceTip):
+  """
+  Inference helper to raise exception on assignment to `rec`.
+  """
+  node_class = astroid.nodes.AssignAttr
+
+  @classmethod
+  def filter(cls, node):
+    if isinstance(node.expr, astroid.nodes.Name) and node.expr.name == 'rec':
+      raise SyntaxError("You can't assign a value to a column with `=`. "
+                        "If you mean to check for equality, use `==` instead.",
+          ('<string>', node.lineno, node.col_offset, ""))
+
+  @classmethod
+  def infer(cls, node, context):
+    raise NotImplementedError()
+
 #----------------------------------------------------------------------
 
 def parse_grist_names(builder):

sandbox/grist/test_codebuilder.py

@@ -76,9 +76,10 @@ class TestCodeBuilder(unittest.TestCase):
                      "# foo($bar=1)\n"
                      "raise SyntaxError('invalid syntax', ('usercode', 1, 5, %s'foo($bar=1)'))"
                      % unicode_prefix)
-    self.assertEqual(make_body('def $bar(): pass'),
-                     "# def $bar(): pass\n"
-                     "raise SyntaxError('invalid syntax', ('usercode', 1, 5, %s'def $bar(): pass'))"
+    self.assertEqual(make_body('def $bar(): return 3'),
+                     "# def $bar(): return 3\n"
+                     "raise SyntaxError('invalid syntax', "
+                     "('usercode', 1, 5, %s'def $bar(): return 3'))"
                      % unicode_prefix)
 
     # If $ is a syntax error, we don't want to turn it into a different syntax error.
@@ -101,20 +102,19 @@ class TestCodeBuilder(unittest.TestCase):
     # Check for reasonable behaviour with non-empty text and no statements.
     self.assertEqual(make_body('# comment'), '# comment\npass')
 
-    self.assertEqual(make_body('rec = 1'), "# rec = 1\n" +
+    self.assertEqual(make_body('rec = 1; rec'), "# rec = 1; rec\n" +
                      "raise SyntaxError('Grist disallows assignment " +
-                     "to the special variable \"rec\"', ('usercode', 1, 1, %s'rec = 1'))"
+                     "to the special variable \"rec\"', ('usercode', 1, 1, %s'rec = 1; rec'))"
                      % unicode_prefix)
-    self.assertEqual(make_body('for rec in []: pass'), "# for rec in []: pass\n" +
+    self.assertEqual(make_body('for rec in []: return rec'), "# for rec in []: return rec\n" +
                      "raise SyntaxError('Grist disallows assignment " +
                      "to the special variable \"rec\"', "
-                     "('usercode', 1, 4, %s'for rec in []: pass'))"
+                     "('usercode', 1, 4, %s'for rec in []: return rec'))"
                      % unicode_prefix)
 
     # some legitimates use of rec
     body = ("""
 foo = rec
-rec.foo = 1
 [rec for x in rec]
 for a in rec:
   t = a
@@ -126,7 +126,6 @@ return rec
     # mostly legitimate use of rec but one failing
     body = ("""
 foo = rec
-rec.foo = 1
 [1 for rec in []]
 for a in rec:
   t = a
@@ -137,9 +136,33 @@ return rec
     self.assertRegex(make_body(body),
                      r"raise SyntaxError\('Grist disallows assignment" +
                      r" to the special variable \"rec\"', "
-                     r"\('usercode', 4, 7, %s'\[1 for rec in \[\]\]'\)\)"
+                     r"\('usercode', 3, 7, %s'\[1 for rec in \[\]\]'\)\)"
+                     % unicode_prefix)
+
+    self.assertEqual(make_body('rec.foo = 1; rec'), "# rec.foo = 1; rec\n" +
+                     "raise SyntaxError(\"You can't assign a value to a column with `=`. "
+                     "If you mean to check for equality, use `==` instead.\", "
+                     "('usercode', 1, 1, %s'rec.foo = 1; rec'))"
+                     % unicode_prefix)
+
+    self.assertEqual(make_body('$foo = 1; rec'), "# $foo = 1; rec\n" +
+                     "raise SyntaxError(\"You can't assign a value to a column with `=`. "
+                     "If you mean to check for equality, use `==` instead.\", "
+                     "('usercode', 1, 1, %s'$foo = 1; rec'))"
+                     % unicode_prefix)
+
+    self.assertEqual(make_body('assert foo'), "# assert foo\n" +
+                     'raise SyntaxError("No `return` statement, '
+                     "and the last line isn't an expression.\", "
+                     "('usercode', 1, 1, %s'assert foo'))"
                      % unicode_prefix)
 
+    self.assertEqual(make_body('foo = 1'), "# foo = 1\n" +
+                     'raise SyntaxError("No `return` statement, '
+                     "and the last line isn't an expression."
+                     " If you want to check for equality, use `==` instead of `=`.\", "
+                     "('usercode', 1, 1, %s'foo = 1'))"
+                     % unicode_prefix)
 
   def test_make_formula_body_unicode(self):
     # Test that we don't fail when strings include unicode characters

sandbox/grist/test_formula_error.py

@@ -35,7 +35,7 @@ else:
         [13, "syntax_err", "Text", True, syntax_err, "", ""],
         [14, "indent_err", "Text", True, indent_err, "", ""],
         [15, "other_err", "Text", True, textwrap.dedent(indent_err), "", ""],
-        [15, "custom_err", "Text", True, "raise Exception('hello')", "", ""],
+        [15, "custom_err", "Text", True, "raise Exception('hello'); return 1", "", ""],
       ]]
     ],
     "DATA": {

sandbox/grist/test_side_effects.py

@@ -40,7 +40,7 @@ class TestSideEffects(test_engine.EngineTestCase):
     # Verify that when a formula fails after a side-effect, the effect is reverted.
     self.load_sample(self.sample)
 
-    formula = 'Schools.lookupOrAddDerived(city="TESTCITY")\nraise Exception("test-error")'
+    formula = 'Schools.lookupOrAddDerived(city="TESTCITY")\nraise Exception("test-error")\nNone'
     out_actions = self.apply_user_action(['AddColumn', 'Address', "A", { 'formula': formula }])
     self.assertPartialOutActions(out_actions, { "stored": [
       ["AddColumn", "Address", "A", {"formula": formula, "isFormula": True, "type": "Any"}],
@@ -70,6 +70,7 @@ class TestSideEffects(test_engine.EngineTestCase):
 Schools.lookupOrAddDerived(city=$city)
 if $amount < 0:
   raise Exception("test-error")
+return None
 '''
     self.add_column('Schools', 'ucity', formula='$city.upper()')
     self.add_column('Address', 'A', formula=formula)