pybricksMicropython: implement appending imports

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

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;
+    });
+}