From 8524b4f791dceea4d36e3fccec1455d3195ea629 Mon Sep 17 00:00:00 2001
From: Alex Hall <alex.mojaki@gmail.com>
Date: Mon, 12 Jul 2021 22:14:23 +0200
Subject: [PATCH] (core) Put user code in linecache so that source lines show
 in tracebacks

Summary: See title. This should improve the user experience, but more importantly it's something I've wanted several times when developing (including just now) so I've been meaning to do this.

Test Plan: Expanded existing unit tests

Reviewers: dsagal, paulfitz

Reviewed By: dsagal

Differential Revision: https://phab.getgrist.com/D2910
---
 sandbox/grist/gencode.py            | 16 ++++++++++++++--
 sandbox/grist/test_formula_error.py | 17 +++++++++++++++--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/sandbox/grist/gencode.py b/sandbox/grist/gencode.py
index 75aeb1b7..12801ad6 100644
--- a/sandbox/grist/gencode.py
+++ b/sandbox/grist/gencode.py
@@ -15,6 +15,7 @@ The schema for grist data is:
     "formula": <opt_string>,
   }
 """
+import linecache
 import re
 import imp
 from collections import OrderedDict
@@ -194,6 +195,17 @@ def _is_special_table(table_id):
 
 def exec_module_text(module_text):
   # pylint: disable=exec-used
-  mod = imp.new_module("usercode")
-  exec(module_text, mod.__dict__)
+  filename = "usercode"
+  mod = imp.new_module(filename)
+
+  # Ensure that source lines show in tracebacks
+  linecache.cache[filename] = (
+    len(module_text),
+    None,
+    [line + '\n' for line in module_text.splitlines()],
+    filename,
+  )
+
+  code_obj = compile(module_text, filename, "exec")
+  exec(code_obj, mod.__dict__)
   return mod
diff --git a/sandbox/grist/test_formula_error.py b/sandbox/grist/test_formula_error.py
index e2605c28..791568f7 100644
--- a/sandbox/grist/test_formula_error.py
+++ b/sandbox/grist/test_formula_error.py
@@ -66,7 +66,14 @@ else:
                               r"TypeError: SQRT\(\) takes 1 positional argument but 2 were given")
 
     self.assertFormulaError(self.engine.get_formula_error('Math', 'built_in_formula', 3),
-                            TypeError, "'int' object is not iterable")
+                            TypeError, "'int' object is not iterable",
+                            textwrap.dedent(
+                              r"""
+                                File "usercode", line \d+, in built_in_formula
+                                  return max\(5\)
+                              TypeError: 'int' object is not iterable
+                              """
+                            ))
 
     self.assertFormulaError(self.engine.get_formula_error('Math', 'syntax_err', 3),
                             SyntaxError, "invalid syntax on line 5 col 9")
@@ -76,7 +83,13 @@ else:
 
     self.assertFormulaError(self.engine.get_formula_error('Math', 'other_err', 3),
                             TypeError, "'int' object is not iterable",
-                            r"line \d+, in other_err")
+                            textwrap.dedent(
+                              r"""
+                                File "usercode", line \d+, in other_err
+                                  if sum\(3, 5\) > 6:
+                              TypeError: 'int' object is not iterable
+                              """
+                            ))
 
     self.assertFormulaError(self.engine.get_formula_error('Math', 'custom_err', 3),
                             Exception, "hello")