Put imported modules in linecache

Author: alexmojaki

backend/main/tests/test_transcript.json

@@ -934,7 +934,7 @@
             "result": [
                 {
                     "color": "red",
-                    "text": "    print('---')\n               ^\nIndentationError: unindent does not match any outer indentation level\nat line 3\n\nAn `IndentationError` occurs when a given line of code is\nnot indented (aligned vertically with other lines) as expected.\n\n"
+                    "text": "    print('---')\n                ^\nIndentationError: unindent does not match any outer indentation level\nat line 3\n\nAn `IndentationError` occurs when a given line of code is\nnot indented (aligned vertically with other lines) as expected.\n\n"
                 },
                 {
                     "color": "red",
@@ -2685,7 +2685,7 @@
             "result": [
                 {
                     "color": "red",
-                    "text": "SyntaxError: expression cannot contain assignment, perhaps you meant \"==\"?\nat line 1\n\nA `SyntaxError` occurs when Python cannot understand your code.\n\n"
+                    "text": "    print(1 + 2 = 3)\n          ^\nSyntaxError: expression cannot contain assignment, perhaps you meant \"==\"?\nat line 1\n\nA `SyntaxError` occurs when Python cannot understand your code.\n\n"
                 },
                 {
                     "color": "red",
@@ -5158,7 +5158,7 @@
             "result": [
                 {
                     "color": "white",
-                    "text": "['__add__', '__class__', '__contains__', '__delattr__', '__delitem__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getitem__', '__gt__', '__hash__', '__iadd__', '__imul__', '__init__', '__init_subclass__', '__iter__', '__le__', '__len__', '__lt__', '__mul__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__reversed__', '__rmul__', '__setattr__', '__setitem__', '__sizeof__', '__str__', '__subclasshook__', 'append', 'clear', 'copy', 'count', 'extend', 'index', 'insert', 'pop', 'remove', 'reverse', 'sort']"
+                    "text": "['__add__', '__class__', '__class_getitem__', '__contains__', '__delattr__', '__delitem__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getitem__', '__gt__', '__hash__', '__iadd__', '__imul__', '__init__', '__init_subclass__', '__iter__', '__le__', '__len__', '__lt__', '__mul__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__reversed__', '__rmul__', '__setattr__', '__setitem__', '__sizeof__', '__str__', '__subclasshook__', 'append', 'clear', 'copy', 'count', 'extend', 'index', 'insert', 'pop', 'remove', 'reverse', 'sort']"
                 },
                 {
                     "color": "white",
@@ -14998,7 +14998,7 @@
             "result": [
                 {
                     "color": "red",
-                    "text": "    is_friend = name == \"Alice\" or\n                                 ^\nSyntaxError: invalid syntax\nat line 1\n\nA `SyntaxError` occurs when Python cannot understand your code.\n\n"
+                    "text": "    is_friend = name == \"Alice\" or\n                                  ^\nSyntaxError: invalid syntax\nat line 1\n\nA `SyntaxError` occurs when Python cannot understand your code.\n\n"
                 },
                 {
                     "color": "red",
@@ -15515,7 +15515,7 @@
             "result": [
                 {
                     "color": "red",
-                    "text": "    \"XOX\n       ^\nSyntaxError: EOL while scanning string literal\nat line 7\n\nA `SyntaxError` occurs when Python cannot understand your code.\n\n"
+                    "text": "    \"XOX\n        ^\nSyntaxError: EOL while scanning string literal\nat line 7\n\nA `SyntaxError` occurs when Python cannot understand your code.\n\n"
                 },
                 {
                     "color": "red",

backend/main/workers/limits.py

@@ -1,5 +1,7 @@
 import inspect
+import linecache
 import os
+import sys
 from functools import lru_cache
 from importlib import import_module
 
@@ -29,15 +31,6 @@ def set_limits():
     import resource
     destroy_dangerous_functions()
 
-    usage = resource.getrusage(resource.RUSAGE_SELF)
-
-    # TODO tests can exceed this time since the process is not restarted, causing failure
-    max_time = int(usage.ru_utime + usage.ru_stime) + 2
-    try:
-        resource.setrlimit(resource.RLIMIT_CPU, (max_time, max_time))
-    except ValueError:
-        pass
-
     patch_cwd()
 
     # Trigger imports before limiting access to files
@@ -59,6 +52,22 @@ def set_limits():
     except NameError as e:
         list(get_suggestions_for_exception(e, e.__traceback__))
 
+    # Put all modules in linecache so that tracebacks work
+    for mod in list(sys.modules.values()):
+        try:
+            linecache.getlines(mod.__file__, mod.__dict__)
+        except Exception:
+            pass
+
+    usage = resource.getrusage(resource.RUSAGE_SELF)
+
+    # TODO tests can exceed this time since the process is not restarted, causing failure
+    max_time = int(usage.ru_utime + usage.ru_stime) + 2
+    try:
+        resource.setrlimit(resource.RLIMIT_CPU, (max_time, max_time))
+    except ValueError:
+        pass
+
     resource.setrlimit(resource.RLIMIT_NOFILE, (0, 0))