(core) Use grist syntactic sugar in AI generated formulas

Summary: Converts `rec.` to `$` in AI generated formulas, and removes redundant `return` at the end.

Test Plan: Expanded unit test.

Reviewers: dsagal

Reviewed By: dsagal

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

sandbox/grist/formula_prompt.py

@@ -4,6 +4,7 @@ import re
 import textwrap
 
 import asttokens
+import asttokens.util
 import six
 
 from column import is_visible_column, BaseReferenceColumn
@@ -230,10 +231,29 @@ def convert_completion(completion):
     if imports:
       result = imports + "\n" + result
 
-    # Check that we still have valid code.
+  # Now convert `rec.` to `$` and remove redundant `return ` at the end.
-    try:
+  try:
-      ast.parse(result)
+    atok = asttokens.ASTTokens(result, parse=True)
-    except SyntaxError:
+  except SyntaxError:
-      return ""
+    # In case the above extraction somehow messed things up
+    return ""
+
+  replacements = []
+  for node in ast.walk(atok.tree):
+    if isinstance(node, ast.Attribute):
+      start, end = atok.get_text_range(node.value)
+      end += 1
+      if result[start:end] == "rec.":
+        replacements.append((start, end, "$"))
+
+  last_stmt = atok.tree.body[-1]
+  if isinstance(last_stmt, ast.Return):
+    start, _ = atok.get_text_range(last_stmt)
+    expected = "return "
+    end = start + len(expected)
+    if result[start:end] == expected:
+      replacements.append((start, end, ""))
+
+  result = asttokens.util.replace(result, replacements)
 
   return result.strip()

sandbox/grist/test_formula_prompt.py

@@ -236,10 +236,12 @@ from x import (
 )
 
 @property
-def foo():
+def foo(rec):
     '''This is a docstring'''
-    x = 5
+    x = f"hello {rec.name} " + rec.name + "!"
-    return 1
+    if rec.bar.spam:
+      return 0
+    return rec.a * rec.b
 ```
 
 Hope you like it!
@@ -251,5 +253,7 @@ from x import (
   z,
 )
 
-x = 5
+x = f"hello {rec.name} " + $name + "!"
-return 1""")
+if $bar.spam:
+  return 0
+$a * $b""")