From 3defb89866692b6c7fc22871b3e3ef784a4a4e81 Mon Sep 17 00:00:00 2001
From: Alex Hall <alex.mojaki@gmail.com>
Date: Fri, 23 Jun 2023 21:33:37 +0200
Subject: [PATCH] (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      | 30 +++++++++++++++++++++++-----
 sandbox/grist/test_formula_prompt.py | 14 ++++++++-----
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/sandbox/grist/formula_prompt.py b/sandbox/grist/formula_prompt.py
index 40b40d20..ba01963f 100644
--- a/sandbox/grist/formula_prompt.py
+++ b/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.
-    try:
-      ast.parse(result)
-    except SyntaxError:
-      return ""
+  # Now convert `rec.` to `$` and remove redundant `return ` at the end.
+  try:
+    atok = asttokens.ASTTokens(result, parse=True)
+  except SyntaxError:
+    # 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()
diff --git a/sandbox/grist/test_formula_prompt.py b/sandbox/grist/test_formula_prompt.py
index cf38e8c4..d03752d4 100644
--- a/sandbox/grist/test_formula_prompt.py
+++ b/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
-    return 1
+    x = f"hello {rec.name} " + rec.name + "!"
+    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
-return 1""")
+x = f"hello {rec.name} " + $name + "!"
+if $bar.spam:
+  return 0
+$a * $b""")