pybricksMicropython: implement appending imports

This implements appending compiled imported modules for the multi-mpy6
file format.

src/fileStorage/index.ts

@@ -29,7 +29,7 @@ export type FileMetadata = Readonly<{
  *
  * IMPORTANT: if this type is changed, we need to modify the database schema to match
  */
-type FileContents = {
+export type FileContents = {
     /** The path of the file in storage. */
     path: string;
     /** The contents of the file. */

src/mpy/sagas.ts

@@ -7,7 +7,8 @@ import { compile as mpyCrossCompileV6 } from '@pybricks/mpy-cross-v6';
 import wasmV6 from '@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm';
 import { call, getContext, put, select, takeEvery } from 'typed-redux-saga/macro';
 import { editorGetValue } from '../editor/sagas';
-import { FileStorageDb } from '../fileStorage';
+import { FileContents, FileStorageDb } from '../fileStorage';
+import { findImportedModules, resolveModule } from '../pybricksMicropython/lib';
 import { RootState } from '../reducers';
 import {
     compile,
@@ -20,6 +21,27 @@ import {
 
 const encoder = new TextEncoder();
 
+/**
+ * Converts JavaScript string to C string.
+ * @param str A string.
+ * @returns Zero-terminated, UTF-8 encoded byte array.
+ */
+function cString(str: string): Uint8Array {
+    return encoder.encode(str + '\x00');
+}
+
+/**
+ * Encodes *value* as a 32-bit unsigned integer in little endian order.
+ * @param value An integer between 0 and 2^32.
+ * @returns A 4-byte array containing the encoded valued.
+ */
+function encodeUInt32LE(value: number): ArrayBuffer {
+    const buf = new ArrayBuffer(4);
+    const view = new DataView(buf);
+    view.setUint32(0, value, true);
+    return buf;
+}
+
 /**
  * Compiles a script to .mpy and dispatches either didCompile on success or
  * didFailToCompile on error.
@@ -77,6 +99,12 @@ function* handleCompile(action: ReturnType<typeof compile>): Generator {
     }
 }
 
+/**
+ * Compiles code into the Pybricks multi-mpy6 file format.
+ *
+ * This includes a __main__ module which is the file currently open in the
+ * editor and any imported modules that can be found in the user file system.
+ */
 function* handleCompileMulti6(): Generator {
     // REVISIT: should we be getting the active file here or have it as an
     // action parameter?
@@ -98,31 +126,79 @@ function* handleCompileMulti6(): Generator {
         return;
     }
 
-    const script = yield* editorGetValue();
+    const mainPy = yield* editorGetValue();
 
-    const result = yield* call(() =>
-        mpyCrossCompileV6(
-            metadata.path ?? '__main__.py',
-            script,
-            undefined,
-            // HACK: testing user agent for jsdom is needed only for getting unit tests to work
-            navigator.userAgent.includes('jsdom') ? undefined : wasmV6,
-        ),
-    );
+    const pyFiles = new Map<string, FileContents>([
+        ['__main__', { path: metadata.path ?? '__main__.py', contents: mainPy }],
+    ]);
 
-    // TODO: add support for imports and append to blob
+    const checkedModules = new Set<string>(['__main__']);
+    const uncheckedScripts = new Array<string>(mainPy);
 
-    if (result.status === 0 && result.mpy) {
-        const sizeBuf = new ArrayBuffer(4);
-        const sizeView = new DataView(sizeBuf);
-        sizeView.setUint32(0, result.mpy.length, true);
+    for (;;) {
+        // parse all unchecked scripts to find imported modules that haven't
+        // been checked yet
 
-        const blob = new Blob([sizeBuf, encoder.encode('__main__\x00'), result.mpy]);
+        const uncheckedModules = new Set<string>();
 
-        yield* put(mpyDidCompileMulti6(blob));
-    } else {
-        yield* put(mpyDidFailToCompileMulti6(result.err));
+        for (const uncheckedScript of uncheckedScripts) {
+            const importedModules = findImportedModules(uncheckedScript);
+
+            for (const m of importedModules) {
+                if (!checkedModules.has(m)) {
+                    uncheckedModules.add(m);
+                }
+            }
+        }
+
+        // all of the scripts have been checked now, so clear the unchecked list
+        uncheckedScripts.length = 0;
+
+        // when no more new modules are found, we are done
+        if (uncheckedModules.size === 0) {
+            break;
+        }
+
+        // try to resolve unchecked modules in the file system
+        for (const m of uncheckedModules) {
+            const file = yield* call(() => resolveModule(db, m));
+
+            // if found, queue the module to be compiled and to be parsed
+            // for additional imports
+            if (file) {
+                pyFiles.set(m, file);
+                uncheckedScripts.push(file.contents);
+            }
+
+            checkedModules.add(m);
+        }
     }
+
+    const blobParts: BlobPart[] = [];
+
+    for (const [m, py] of pyFiles) {
+        const result = yield* call(() =>
+            mpyCrossCompileV6(
+                py.path,
+                py.contents,
+                undefined,
+                // HACK: testing user agent for jsdom is needed only for getting unit tests to work
+                navigator.userAgent.includes('jsdom') ? undefined : wasmV6,
+            ),
+        );
+
+        if (result.status !== 0 || !result.mpy) {
+            yield* put(mpyDidFailToCompileMulti6(result.err));
+            return;
+        }
+
+        // each file is encoded as the size, module name, and mpy binary
+        blobParts.push(encodeUInt32LE(result.mpy.length));
+        blobParts.push(cString(m));
+        blobParts.push(result.mpy);
+    }
+
+    yield* put(mpyDidCompileMulti6(new Blob(blobParts)));
 }
 
 export default function* (): Generator {

src/pybricksMicropython/lib.test.ts

@@ -3,6 +3,7 @@
 
 import {
     FileNameValidationResult,
+    findImportedModules,
     pythonFileExtension,
     pythonFileExtensionRegex,
     validateFileName,
@@ -76,3 +77,46 @@ describe('validateFileName', () => {
         );
     });
 });
+
+test('findImportedModules', async () => {
+    const script = `
+import a
+import b, c
+import d.d
+import e.e as e
+import f.f as f, g
+from h import x
+from h import x as y
+from i import (x, y)
+from i import (x as y, z)
+from j import *
+from . import x
+from . import x as y
+from .r import x
+from ..r import x
+from ...r import x
+from ....r import x
+`;
+
+    const modules = findImportedModules(script);
+
+    expect(modules).toEqual(
+        new Set([
+            'a',
+            'b',
+            'c',
+            'd.d',
+            'e.e',
+            'f.f',
+            'g',
+            'h',
+            'i',
+            'j',
+            '.',
+            '.r',
+            '..r',
+            '...r',
+            '....r',
+        ]),
+    );
+});

src/pybricksMicropython/lib.ts

@@ -1,6 +1,9 @@
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2022 The Pybricks Authors
 
+import { parse, walk } from 'python-ast';
+import type { FileContents, FileStorageDb } from '../fileStorage';
+
 /** The Python file extension ('.py') */
 export const pythonFileExtension = '.py';
 
@@ -67,3 +70,78 @@ export function validateFileName(
 
     return FileNameValidationResult.IsOk;
 }
+
+/**
+ * Finds modules imported by a Python script.
+ *
+ * @param py A Python Script.
+ * @returns A list of the names of modules imported by this file.
+ */
+export function findImportedModules(py: string): ReadonlySet<string> {
+    const modules = new Set<string>();
+    const tree = parse(py);
+
+    // find all import statements in the syntax tree and collect imported modules
+    walk(
+        {
+            enterImport_stmt: (ctx) => {
+                // import statements have two forms:
+                // import_stmt: import_name | import_from;
+
+                // import_name: 'import' dotted_as_names
+                const name = ctx.import_name();
+
+                if (name) {
+                    // dotted_as_names: dotted_as_name (',' dotted_as_name)*;
+                    // dotted_as_name: dotted_name ('as' NAME)?;
+                    for (const dottedAsName of name
+                        .dotted_as_names()
+                        .dotted_as_name()) {
+                        // dotted_name: NAME ('.' NAME)*;
+                        modules.add(dottedAsName.dotted_name().text);
+                    }
+                }
+
+                // import_from: ('from' (('.' | '...')* dotted_name | ('.' | '...')+) 'import' ('*' | '(' import_as_names ')' | import_as_names ));
+                const from = ctx.import_from();
+
+                if (from) {
+                    // the leading dots aren't included in dotted_name, so
+                    // we need to collect them separately
+                    const leadingDots = from
+                        .DOT()
+                        .concat(from.ELLIPSIS())
+                        .map((n) => n.symbol.text)
+                        .join('');
+
+                    // dotted_name: NAME ('.' NAME)*;
+                    const dottedName = from.dotted_name()?.text ?? '';
+
+                    modules.add(leadingDots + dottedName);
+                }
+            },
+        },
+        tree,
+    );
+
+    return modules;
+}
+
+export async function resolveModule(
+    db: FileStorageDb,
+    module: string,
+): Promise<FileContents | undefined> {
+    const modulePath = module.replace(/\./g, '/') + '.py';
+
+    return await db.transaction('r', db.metadata, db._contents, async () => {
+        const match = await db.metadata.where('path').equals(modulePath).first();
+
+        if (!match) {
+            return undefined;
+        }
+
+        const file = await db._contents.get(match.path);
+
+        return file;
+    });
+}