Refactor pyexec (#1318)

This is some refactoring I did on the way towards resolving pyscript#1313.
I added a new _run_pyscript Python function which executes the code
inside a context manager that sets the display target. We can then
return a JS object wrapper directly from Python.

I moved the "installation" of the pyscript module to loadInterpreter,
and pyimport pyscript_py there and give it a type. This avoids a bunch
of creating and deleting of proxies for pyscript_py and allows us to
give it a type once and for all.

I also did some minor logic cleanup in a few places.

pyscriptjs/src/interpreter_client.ts

@@ -45,13 +45,23 @@ export class InterpreterClient extends Object {
     }
 
     /**
-     * delegates the code to be run to the underlying interface of
-     * the remote interpreter.
-     * Python exceptions are turned into JS exceptions.
-     * */
+     * Run user Python code. See also the _run_pyscript docstring.
+     *
+     * The result is wrapped in an object to avoid accidentally awaiting a
+     * Python Task or Future returned as the result of the computation.
+     *
+     * @param code the code to run
+     * @param id The id for the default display target (or undefined if no
+     * default display target).
+     * @returns Either:
+     * 1. An Object of the form {result: the_result} if the result is
+     *    serializable (or transferable), or
+     * 2. a Synclink Proxy wrapping an object of this if the result is not
+     *    serializable.
+     */
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    async run(code: string): Promise<{ result: any }> {
-        return await this._remote.run(code);
+    async run(code: string, id?: string): Promise<{ result: any }> {
+        return this._remote.pyscript_py._run_pyscript(code, id);
     }
 
     /**

pyscriptjs/src/main.ts

@@ -21,14 +21,6 @@ import { RemoteInterpreter } from './remote_interpreter';
 import { robustFetch } from './fetch';
 import * as Synclink from 'synclink';
 
-// pyscript_package is injected from src/python by bundlePyscriptPythonPlugin in
-// esbuild.js
-
-// @ts-ignore
-import python_package from 'pyscript_python_package.esbuild_injected.json';
-
-declare const python_package: { dirs: string[]; files: [string, string] };
-
 const logger = getLogger('pyscript/main');
 
 /**
@@ -271,17 +263,6 @@ export class PyScriptApp {
         // XXX: maybe the following calls could be parallelized, instead of
         // await()ing immediately. For now I'm using await to be 100%
         // compatible with the old behavior.
-        logger.info('importing pyscript');
-
-        // Write pyscript package into file system
-        for (const dir of python_package.dirs) {
-            await interpreter._remote.FS.mkdir('/home/pyodide/' + dir);
-        }
-        for (const [path, value] of python_package.files) {
-            await interpreter._remote.FS.writeFile('/home/pyodide/' + path, value);
-        }
-        //Refresh the module cache so Python consistently finds pyscript module
-        await interpreter._remote.invalidate_module_path_cache();
 
         // inject `define_custom_element` and showWarning it into the PyScript
         // module scope
@@ -290,11 +271,7 @@ export class PyScriptApp {
         // await interpreter._remote.setHandler('showWarning', Synclink.proxy(showWarning));
         interpreter._unwrapped_remote.setHandler('define_custom_element', define_custom_element);
         interpreter._unwrapped_remote.setHandler('showWarning', showWarning);
-        const pyscript_module = (await interpreter.pyimport('pyscript')) as Synclink.Remote<
-            PyProxy & { _set_version_info(string): void }
-        >;
-        await pyscript_module._set_version_info(version);
-        await pyscript_module.destroy();
+        await interpreter._remote.pyscript_py._set_version_info(version);
 
         // import some carefully selected names into the global namespace
         await interpreter.run(`

pyscriptjs/src/pyexec.ts

@@ -2,8 +2,7 @@ import { getLogger } from './logger';
 import { ensureUniqueId } from './utils';
 import { UserError, ErrorCode } from './exceptions';
 import { InterpreterClient } from './interpreter_client';
-import type { PyProxyCallable, PyProxy } from 'pyodide';
-import type { Remote } from 'synclink';
+import type { PyProxyCallable } from 'pyodide';
 
 const logger = getLogger('pyexec');
 
@@ -13,39 +12,30 @@ export async function pyExec(
     outElem: HTMLElement,
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
 ): Promise<{ result: any }> {
-    const pyscript_py = (await interpreter.pyimport('pyscript')) as Remote<
-        PyProxy & {
-            set_current_display_target(id: string): void;
-            uses_top_level_await(code: string): boolean;
-        }
-    >;
     ensureUniqueId(outElem);
-    await pyscript_py.set_current_display_target(outElem.id);
+    if (await interpreter._remote.pyscript_py.uses_top_level_await(pysrc)) {
+        const err = new UserError(
+            ErrorCode.TOP_LEVEL_AWAIT,
+            'The use of top-level "await", "async for", and ' +
+                '"async with" has been removed.' +
+                '\nPlease write a coroutine containing ' +
+                'your code and schedule it using asyncio.ensure_future() or similar.' +
+                '\nSee https://docs.pyscript.net/latest/guides/asyncio.html for more information.',
+        );
+        displayPyException(err, outElem);
+        return { result: undefined };
+    }
+
     try {
-        try {
-            if (await pyscript_py.uses_top_level_await(pysrc)) {
-                throw new UserError(
-                    ErrorCode.TOP_LEVEL_AWAIT,
-                    'The use of top-level "await", "async for", and ' +
-                        '"async with" has been removed.' +
-                        '\nPlease write a coroutine containing ' +
-                        'your code and schedule it using asyncio.ensure_future() or similar.' +
-                        '\nSee https://docs.pyscript.net/latest/guides/asyncio.html for more information.',
-                );
-            }
-            return await interpreter.run(pysrc);
-        } catch (e) {
-            const err = e as Error;
-            // XXX: currently we display exceptions in the same position as
-            // the output. But we probably need a better way to do that,
-            // e.g. allowing plugins to intercept exceptions and display them
-            // in a configurable way.
-            displayPyException(err, outElem);
-            return { result: undefined };
-        }
-    } finally {
-        await pyscript_py.set_current_display_target(undefined);
-        await pyscript_py.destroy();
+        return await interpreter.run(pysrc, outElem.id);
+    } catch (e) {
+        const err = e as Error;
+        // XXX: currently we display exceptions in the same position as
+        // the output. But we probably need a better way to do that,
+        // e.g. allowing plugins to intercept exceptions and display them
+        // in a configurable way.
+        displayPyException(err, outElem);
+        return { result: undefined };
     }
 }
 

pyscriptjs/src/python/pyscript/__init__.py

@@ -6,14 +6,17 @@ import io
 import re
 import time
 from collections import namedtuple
+from contextlib import contextmanager
 from textwrap import dedent
 
 import js
 
 try:
-    from pyodide.ffi import create_proxy
+    from pyodide.code import eval_code
+    from pyodide.ffi import JsProxy, create_proxy
 except ImportError:
-    from pyodide import create_proxy
+    from pyodide import JsProxy, create_proxy, eval_code
+
 
 loop = asyncio.get_event_loop()
 
@@ -188,8 +191,13 @@ def write(element_id, value, append=False, exec_id=0):
     )
 
 
-def set_current_display_target(target_id):
+@contextmanager
+def _display_target(target_id):
     get_current_display_target._id = target_id
+    try:
+        yield
+    finally:
+        get_current_display_target._id = None
 
 
 def get_current_display_target():
@@ -200,17 +208,14 @@ get_current_display_target._id = None
 
 
 def display(*values, target=None, append=True):
-    default_target = get_current_display_target()
-    if default_target is None and target is None:
+    if target is None:
+        target = get_current_display_target()
+    if target is None:
         raise Exception(
             "Implicit target not allowed here. Please use display(..., target=...)"
         )
-    if target is not None:
-        for v in values:
-            Element(target).write(v, append=append)
-    else:
-        for v in values:
-            Element(default_target).write(v, append=append)
+    for v in values:
+        Element(target).write(v, append=append)
 
 
 class Element:
@@ -702,3 +707,32 @@ def _install_deprecated_globals_2022_12_1(ns):
         "Please use <code>pyscript</code> instead."
     )
     ns["PyScript"] = DeprecatedGlobal("PyScript", PyScript, message)
+
+
+def _run_pyscript(code: str, id: str = None) -> JsProxy:
+    """Execute user code inside context managers.
+
+    Uses the __main__ global namespace.
+
+    The output is wrapped inside a JavaScript object since an object is not
+    thenable. This is so we do not accidentally `await` the result of the python
+    execution, even if it's awaitable (Future, Task, etc.)
+
+    Parameters
+    ----------
+    code :
+       The code to run
+
+    id :
+       The id for the default display target (or None if no default display target).
+
+    Returns
+    -------
+        A Js Object of the form {result: the_result}
+    """
+    import __main__
+
+    with _display_target(id):
+        result = eval_code(code, globals=__main__.__dict__)
+
+    return js.Object.new(result=result)

pyscriptjs/src/python_package.ts

@@ -0,0 +1,9 @@
+// This file exists because I can only convince jest to mock real file system
+// files, not fake modules. This confuses me because jest.mock has a virtual
+// option for mocking things that don't live in the file system but it doesn't
+// seem to work.
+
+// @ts-ignore
+import python_package from 'pyscript_python_package.esbuild_injected.json';
+declare const python_package: { dirs: string[]; files: [string, string][] };
+export { python_package };

pyscriptjs/src/remote_interpreter.ts

@@ -7,6 +7,8 @@ import type { loadPyodide as loadPyodideDeclaration, PyodideInterface, PyProxy,
 import type { ProxyMarked } from 'synclink';
 import * as Synclink from 'synclink';
 
+import { python_package } from './python_package';
+
 declare const loadPyodide: typeof loadPyodideDeclaration;
 const logger = getLogger('pyscript/pyodide');
 
@@ -30,6 +32,13 @@ type PATHInterface = {
     dirname(path: string): string;
 } & ProxyMarked;
 
+type PyScriptPyModule = ProxyMarked & {
+    _set_version_info(ver: string): void;
+    uses_top_level_await(code: string): boolean;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    _run_pyscript(code: string, display_target_id?: string): { result: any };
+};
+
 /*
 RemoteInterpreter class is responsible to process requests from the
 `InterpreterClient` class -- these can be requests for installation of
@@ -51,6 +60,7 @@ export class RemoteInterpreter extends Object {
     FS: FSInterface;
     PATH: PATHInterface;
     PATH_FS: PATHFSInterface;
+    pyscript_py: PyScriptPyModule;
 
     globals: PyProxyDict & ProxyMarked;
     // TODO: Remove this once `runtimes` is removed!
@@ -106,40 +116,28 @@ export class RemoteInterpreter extends Object {
         // TODO: Remove this once `runtimes` is removed!
         this.interpreter = this.interface;
 
+        // Write pyscript package into file system
+        for (const dir of python_package.dirs) {
+            this.FS.mkdir('/home/pyodide/' + dir);
+        }
+        for (const [path, value] of python_package.files) {
+            this.FS.writeFile('/home/pyodide/' + path, value);
+        }
+        //Refresh the module cache so Python consistently finds pyscript module
+        this.invalidate_module_path_cache();
+
         this.globals = Synclink.proxy(this.interface.globals as PyProxyDict);
+        logger.info('importing pyscript');
+        this.pyscript_py = Synclink.proxy(this.interface.pyimport('pyscript')) as PyProxy & typeof this.pyscript_py;
 
         if (config.packages) {
             logger.info('Found packages in configuration to install. Loading micropip...');
             await this.loadPackage('micropip');
         }
         logger.info('pyodide loaded and initialized');
-        await this.run('print("Python initialization complete")');
+        this.pyscript_py._run_pyscript('print("Python initialization complete")');
     }
 
-    /* eslint-disable */
-    async run(code: string): Promise<{ result: any }> {
-        /**
-         * eslint wants `await` keyword to be used i.e.
-         * { result: await this.interface.runPython(code) }
-         * However, `await` is not a no-op (no-operation) i.e.
-         * `await 42` is NOT the same as `42` i.e. if the awaited
-         * thing is not a promise, it is wrapped inside a promise and
-         * that promise is awaited. Thus, it changes the execution order.
-         * See https://stackoverflow.com/questions/55262996/does-awaiting-a-non-promise-have-any-detectable-effect
-         * Thus, `eslint` is disabled for this block / snippet.
-         */
-
-        /**
-         * The output of `runPython` is wrapped inside an object
-         * since an object is not thennable and avoids return of
-         * a coroutine directly. This is so we do not `await` the results
-         * of the underlying python execution, even if it's an
-         * awaitable object (Future, Task, etc.)
-         */
-        return { result: this.interface.runPython(code) };
-    }
-    /* eslint-enable */
-
     /**
      * delegates the registration of JS modules to
      * the underlying interface.