safe_exec can load modules for you.

common/lib/codejail/codejail/safe_exec.py

@@ -2,6 +2,8 @@
 
 import json
 
+from .lazymod import LazyModule
+
 def straw(v):
     return json.loads(json.dumps(jsonable_dict(v)))
 
@@ -27,6 +29,13 @@ def safe_exec(code, globals_dict, locals_dict=None, future_division=False, assum
     else:
         l_dict = straw(locals_dict)
 
+    for modname in assumed_imports or ():
+        if isinstance(modname, tuple):
+            name, modname = modname
+        else:
+            name = modname
+        g_dict[name] = LazyModule(modname)
+
     exec code in g_dict, l_dict
 
     globals_dict.update(straw(g_dict))