fixes #28

Author: jph00

01_funccall.ipynb

@@ -1368,12 +1368,12 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "84cadbec-0b34-4669-bfc9-6bf790816035",
+   "id": "1574585f",
    "metadata": {},
    "outputs": [],
    "source": [
     "#| exports\n",
-    "def _run(code:str ):\n",
+    "def _run(code:str, glb:dict=None, loc:dict=None):\n",
     "    \"Run `code`, returning final expression (similar to IPython)\"\n",
     "    tree = ast.parse(code)\n",
     "    last_node = tree.body[-1] if tree.body else None\n",
@@ -1385,13 +1385,13 @@
     "        tree.body[-1] = _copy_loc(assign_node, last_node)\n",
     "\n",
     "    compiled_code = compile(tree, filename='<ast>', mode='exec')\n",
-    "    namespace = {}\n",
+    "    glb = glb or {}\n",
     "    stdout_buffer = io.StringIO()\n",
     "    saved_stdout = sys.stdout\n",
     "    sys.stdout = stdout_buffer\n",
-    "    try: exec(compiled_code, namespace)\n",
+    "    try: exec(compiled_code, glb, loc)\n",
     "    finally: sys.stdout = saved_stdout\n",
-    "    _result = namespace.get('_result', None)\n",
+    "    _result = glb.get('_result', None)\n",
     "    if _result is not None: return _result\n",
     "    return stdout_buffer.getvalue().strip()"
    ]
@@ -1462,15 +1462,19 @@
    "outputs": [],
    "source": [
     "#| exports\n",
-    "def python(code, # Code to execute\n",
-    "           timeout=5 # Maximum run time in seconds before a `TimeoutError` is raised\n",
+    "def python(code:str, # Code to execute\n",
+    "           glb:Optional[dict]=None, # Globals namespace\n",
+    "           loc:Optional[dict]=None, # Locals namespace\n",
+    "           timeout:int=3600 # Maximum run time in seconds before a `TimeoutError` is raised\n",
     "          ): # Result of last node, if it's an expression, or `None` otherwise\n",
     "    \"\"\"Executes python `code` with `timeout` and returning final expression (similar to IPython).\n",
     "    Raised exceptions are returned as a string, with a stack trace.\"\"\"\n",
     "    def handler(*args): raise TimeoutError()\n",
+    "    if glb is None: glb = inspect.currentframe().f_back.f_globals\n",
+    "    if loc is None: loc=glb\n",
     "    signal.signal(signal.SIGALRM, handler)\n",
     "    signal.alarm(timeout)\n",
-    "    try: return _run(code)\n",
+    "    try: return _run(code, glb, loc)\n",
     "    except Exception as e: return traceback.format_exc()\n",
     "    finally: signal.alarm(0)"
    ]
@@ -1512,18 +1516,98 @@
    "id": "6c629442",
    "metadata": {},
    "source": [
-    "If the code takes longer than `timeout` then it raises a `TimeoutError`."
+    "If the code takes longer than `timeout` then it returns an error string."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "fcb472b3",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Traceback (most recent call last):\n",
+      "  File \"/var/folders/51/b2_szf2945n072c0vj2cyty40000gn/T/ipykernel_12053/775515280.py\", line 12, in python\n",
+      "    try: return _run(code, glb, loc)\n",
+      "                ^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/var/folders/51/b2_szf2945n072c0vj2cyty40000gn/T/ipykernel_12053/1858893181.py\", line 18, in _run\n",
+      "    try: exec(compiled_code, glb, loc)\n",
+      "         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"<ast>\", line 1, in <module>\n",
+      "  File \"/var/folders/51/b2_szf2945n072c0vj2cyty40000gn/T/ipykernel_12053/775515280.py\", line 9, in handler\n",
+      "    def handler(*args): raise TimeoutError()\n",
+      "                        ^^^^^^^^^^^^^^^^^^^^\n",
+      "TimeoutError\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(python('import time; time.sleep(10)', timeout=1))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d45684c1",
+   "metadata": {},
+   "source": [
+    "By default the caller's global namespace is used."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72dfe290",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "1"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "python(\"a=1\")\n",
+    "a"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bf48557c",
+   "metadata": {},
+   "source": [
+    "Pass a different `glb` if needed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "55fb5613",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(1, 3)"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "try: python('import time; time.sleep(10)', timeout=1)\n",
-    "except TimeoutError: print('Timed out')"
+    "glb = {}\n",
+    "python(\"a=3\", glb)\n",
+    "a, glb['a']"
    ]
   },
   {

toolslm/funccall.py

@@ -141,7 +141,7 @@ def _copy_loc(new, orig):
     return new
 
 # %% ../01_funccall.ipynb 69
-def _run(code:str ):
+def _run(code:str, glb:dict=None, loc:dict=None):
     "Run `code`, returning final expression (similar to IPython)"
     tree = ast.parse(code)
     last_node = tree.body[-1] if tree.body else None
@@ -153,30 +153,34 @@ def _run(code:str ):
         tree.body[-1] = _copy_loc(assign_node, last_node)
 
     compiled_code = compile(tree, filename='<ast>', mode='exec')
-    namespace = {}
+    glb = glb or {}
     stdout_buffer = io.StringIO()
     saved_stdout = sys.stdout
     sys.stdout = stdout_buffer
-    try: exec(compiled_code, namespace)
+    try: exec(compiled_code, glb, loc)
     finally: sys.stdout = saved_stdout
-    _result = namespace.get('_result', None)
+    _result = glb.get('_result', None)
     if _result is not None: return _result
     return stdout_buffer.getvalue().strip()
 
 # %% ../01_funccall.ipynb 74
-def python(code, # Code to execute
-           timeout=5 # Maximum run time in seconds before a `TimeoutError` is raised
+def python(code:str, # Code to execute
+           glb:Optional[dict]=None, # Globals namespace
+           loc:Optional[dict]=None, # Locals namespace
+           timeout:int=3600 # Maximum run time in seconds before a `TimeoutError` is raised
           ): # Result of last node, if it's an expression, or `None` otherwise
     """Executes python `code` with `timeout` and returning final expression (similar to IPython).
     Raised exceptions are returned as a string, with a stack trace."""
     def handler(*args): raise TimeoutError()
+    if glb is None: glb = inspect.currentframe().f_back.f_globals
+    if loc is None: loc=glb
     signal.signal(signal.SIGALRM, handler)
     signal.alarm(timeout)
-    try: return _run(code)
+    try: return _run(code, glb, loc)
     except Exception as e: return traceback.format_exc()
     finally: signal.alarm(0)
 
-# %% ../01_funccall.ipynb 81
+# %% ../01_funccall.ipynb 85
 def mk_ns(*funcs_or_objs):
     merged = {}
     for o in funcs_or_objs:
@@ -185,14 +189,14 @@ def mk_ns(*funcs_or_objs):
         if callable(o) and hasattr(o, '__name__'): merged |= {o.__name__: o}
     return merged
 
-# %% ../01_funccall.ipynb 90
+# %% ../01_funccall.ipynb 94
 def call_func(fc_name, fc_inputs, ns):
     "Call the function `fc_name` with the given `fc_inputs` using namespace `ns`."
     if not isinstance(ns, abc.Mapping): ns = mk_ns(*ns)
     func = ns[fc_name]
     return func(**fc_inputs)
 
-# %% ../01_funccall.ipynb 97
+# %% ../01_funccall.ipynb 101
 async def call_func_async(fc_name, fc_inputs, ns):
     "Awaits the function `fc_name` with the given `fc_inputs` using namespace `ns`."
     if not isinstance(ns, abc.Mapping): ns = mk_ns(*ns)