From 0430d340024f07b9add0b66f7c8361dc5367f16e Mon Sep 17 00:00:00 2001 From: David Lechner Date: Mon, 7 Sep 2026 13:24:22 -0500 Subject: [PATCH] mpy/sagas: Get imported modules from compiled .mpy. We were parsing Python in JavaScript to find imports, using a jison parser that allocated memory proportional to lines times file size, so large programs could exhaust browser memory. The grammar also rejected f-strings containing the same quote character and decimal literals with a leading zero, and since parse errors were swallowed, any such program silently downloaded with no modules at all. Read the imports back out of the .mpy instead. Each module is already compiled by mpy-cross, so this removes a pass rather than adding one, and MicroPython's own compiler decides what an import is. Fixes: https://github.com/pybricks/support/issues/1804 Fixes: https://github.com/pybricks/support/issues/1954 Fixes: https://github.com/pybricks/support/issues/1981 Assisted-by: Claude Opus 5 --- package.json | 1 - src/mpy/mpyImports.test.ts | 277 ++++++++++++++++++++++++ src/mpy/mpyImports.ts | 322 ++++++++++++++++++++++++++++ src/mpy/sagas.test.ts | 204 +++++++++++++++--- src/mpy/sagas.ts | 113 +++++----- src/mpy/staticQstrs.ts | 187 ++++++++++++++++ src/mpy/test-utils.ts | 49 +++++ src/pybricksMicropython/lib.test.ts | 70 ------ src/pybricksMicropython/lib.ts | 39 ---- yarn.lock | 10 - 10 files changed, 1073 insertions(+), 199 deletions(-) create mode 100644 src/mpy/mpyImports.test.ts create mode 100644 src/mpy/mpyImports.ts create mode 100644 src/mpy/staticQstrs.ts create mode 100644 src/mpy/test-utils.ts diff --git a/package.json b/package.json index 4c2930aa6..37a69bc38 100644 --- a/package.json +++ b/package.json @@ -17,7 +17,6 @@ "@pybricks/jedi": "1.17.0", "@pybricks/mpy-cross-v5": "^2.0.0", "@pybricks/mpy-cross-v6": "^2.0.0", - "@pybricks/python-program-analysis": "^2.0.0", "@pyodide/webpack-plugin": "^1.3.2", "@reduxjs/toolkit": "^1.9.7", "@shopify/react-i18n": "^7.13.1", diff --git a/src/mpy/mpyImports.test.ts b/src/mpy/mpyImports.test.ts new file mode 100644 index 000000000..b6b96ffd3 --- /dev/null +++ b/src/mpy/mpyImports.test.ts @@ -0,0 +1,277 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 The Pybricks Authors + +import { compile as mpyCrossCompileV6 } from '@pybricks/mpy-cross-v6'; +import { MpyFormatError, findImportedModules } from './mpyImports'; +import { mockMpyCrossWasmPath, mpyCrossV6WasmPath } from './test-utils'; + +beforeEach(() => { + mockMpyCrossWasmPath(); +}); + +afterEach(() => { + jest.clearAllMocks(); +}); + +async function compileToMpy(script: string): Promise { + const result = await mpyCrossCompileV6( + 'test.py', + script, + undefined, + mpyCrossV6WasmPath, + ); + + expect(result.err).toEqual([]); + expect(result.status).toBe(0); + expect(result.mpy).toBeDefined(); + + return result.mpy as Uint8Array; +} + +async function findImportsInScript(script: string): Promise> { + return findImportedModules(await compileToMpy(script)); +} + +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 + +# import q +# from q import q +""" +import q +from q import q +""" +''' +import q +from q import q +''' +`; + + const modules = await findImportsInScript(script); + + expect(modules).toEqual( + new Set([ + 'a', + 'b', + 'c', + 'd.d', + 'e.e', + 'f.f', + 'g', + 'h', + 'i', + 'j', + '.', + '.r', + '..r', + '...r', + '....r', + ]), + ); +}); + +test('https://github.com/pybricks/support/issues/873 regression', async () => { + const script = ` +from my_module import data + +async def hello(): + print("hello") + +print(data) +`; + + const modules = await findImportsInScript(script); + + expect(modules).toEqual(new Set(['my_module'])); +}); + +test('https://github.com/pybricks/support/issues/1954 regression', async () => { + // an f-string containing the same kind of quote that delimits it (PEP 701) + const script = ` +import mission_run_1 +missions = [("1", "irrelevant_name", "first"), ("2", "stop", "stop"),] +print(f"hasattr f-string {hasattr(missions[0][1], "run")}") +`; + + const modules = await findImportsInScript(script); + + expect(modules).toEqual(new Set(['mission_run_1'])); +}); + +test('https://github.com/pybricks/support/issues/1981 regression', async () => { + // MicroPython accepts a decimal literal with a leading zero + const script = ` +import my_lib +print(050) +`; + + const modules = await findImportsInScript(script); + + expect(modules).toEqual(new Set(['my_lib'])); +}); + +test( + 'https://github.com/pybricks/support/issues/1804 regression', + async () => { + // A large program used to make the old Python parser allocate memory + // proportional to (lines * file size), which crashed the browser tab. + const line = + 'data = b"132456789132456789132456789132456789132456789132456789132456789"'; + const script = ['import my_lib', ...new Array(7000).fill(line), ''].join('\n'); + + expect(script.length).toBeGreaterThan(400000); + + const modules = await findImportsInScript(script); + + expect(modules).toEqual(new Set(['my_lib'])); + }, + 30 * 1000, +); + +test('imports are found in every scope', async () => { + const script = ` +import at_module_level + +def func(): + import in_func + +class Cls: + import in_class + + def method(self): + import in_method + +async def async_func(): + import in_async_func + +def outer(): + def inner(): + import deeply_nested + +if unknown: + import in_if +else: + import in_else + +while unknown: + import in_while + +for i in range(3): + import in_for + +with open("f") as f: + import in_with + +try: + import in_try +except ImportError: + import in_except +finally: + import in_finally + +lam = lambda: [__import__("in_comprehension") for i in range(1)] +`; + + const modules = await findImportsInScript(script); + + expect(modules).toEqual( + new Set([ + 'at_module_level', + 'in_func', + 'in_class', + 'in_method', + 'in_async_func', + 'deeply_nested', + 'in_if', + 'in_else', + 'in_while', + 'in_for', + 'in_with', + 'in_try', + 'in_except', + 'in_finally', + ]), + ); +}); + +test('module names that are static qstrs are found', async () => { + // Names at or below QSTR_LAST_STATIC are stored as an index instead of a string. + // 'main' is one of them and main.py is a very common user program name. + const modules = await findImportsInScript( + 'import main\nimport time\nimport sys\nfrom math import pi\n', + ); + + expect(modules).toEqual(new Set(['main', 'time', 'sys', 'math'])); +}); + +test('a program with no imports gives an empty set', async () => { + const modules = await findImportsInScript('print("hello!")\n'); + + expect(modules).toEqual(new Set()); +}); + +describe('bad .mpy files are rejected', () => { + // These guard against a future mpy-cross update changing the file format. If any of + // them start failing, findImportedModules() needs to be updated to match. + + test('the expected header is produced', async () => { + const mpy = await compileToMpy('print("hello!")\n'); + + expect(mpy[0]).toBe('M'.charCodeAt(0)); + expect(mpy[1]).toBe(6); // ABI version + expect(mpy[2]).toBe(0); // feature flags (no native code) + expect(mpy[3]).toBe(31); // small int bits + }); + + test('not a .mpy file', () => { + expect(() => findImportedModules(new Uint8Array([1, 2, 3, 4]))).toThrow( + MpyFormatError, + ); + }); + + test('empty file', () => { + expect(() => findImportedModules(new Uint8Array())).toThrow(MpyFormatError); + }); + + test('unsupported abi version', async () => { + const mpy = await compileToMpy('print("hello!")\n'); + mpy[1] = 7; + + expect(() => findImportedModules(mpy)).toThrow( + 'unsupported .mpy ABI version: 7', + ); + }); + + test('contains native code', async () => { + const mpy = await compileToMpy('print("hello!")\n'); + mpy[2] = 1 << 2; + + expect(() => findImportedModules(mpy)).toThrow( + '.mpy file contains native code', + ); + }); + + test('truncated file', async () => { + const mpy = await compileToMpy('import my_lib\nprint("hello!")\n'); + + expect(() => findImportedModules(mpy.subarray(0, mpy.length - 5))).toThrow( + MpyFormatError, + ); + }); +}); diff --git a/src/mpy/mpyImports.ts b/src/mpy/mpyImports.ts new file mode 100644 index 000000000..560affaba --- /dev/null +++ b/src/mpy/mpyImports.ts @@ -0,0 +1,322 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 The Pybricks Authors + +import { staticQstrs } from './staticQstrs'; + +/** + * Reads the names of the modules imported by a program from its compiled .mpy file. + * + * Doing this instead of parsing the Python source in JavaScript means we get exactly + * the same answer MicroPython itself would get, since it is MicroPython's own compiler + * that produced the bytecode we are reading. + * + * The format is documented by `py/persistentcode.c` and `tools/mpy-tool.py` in + * MicroPython. Everything here matches MicroPython v1.19.1 (which + * `@pybricks/mpy-cross-v6` is built from) through pybricks-micropython master, and the + * `MPY_SUB_VERSION` mechanism exists so that bytecode-only .mpy files stay compatible + * across all of ABI v6. + */ + +/** The .mpy ABI version this code knows how to read. */ +const mpyAbiVersion = 6; + +// Opcodes from py/bc0.h that we need to recognize. +const mpBcLoadConstSmallInt = 0x22; +const mpBcImportName = 0x1b; +const mpBcLoadConstSmallIntMulti = 0x70; +const mpBcLoadConstSmallIntMultiNum = 64; +const mpBcLoadConstSmallIntMultiExcess = 16; + +/** Opcodes with `opcode & this === 0` are followed by one extra byte. */ +const mpBcMaskExtraByte = 0x9e; + +// Operand encodings from py/bc0.h. The packed table maps the high nibble of an opcode +// to one of these; see MP_BC_FORMAT() and mp_opcode_decode(). +const mpBcFormatTable = 0x3a4; +const mpBcFormatQstr = 1; +const mpBcFormatVarUint = 2; +const mpBcFormatOffset = 3; + +/** Object types from the MP_PERSISTENT_OBJ_* enum in py/persistentcode.h. */ +enum PersistentObj { + FunTable = 0, + None = 1, + False = 2, + True = 3, + Ellipsis = 4, + Str = 5, + Bytes = 6, + Int = 7, + Float = 8, + Complex = 9, + Tuple = 10, +} + +/** Error raised when a .mpy file cannot be read. */ +export class MpyFormatError extends Error {} + +/** + * Finds the modules imported by a compiled program. + * + * Relative imports are returned with their leading dots, e.g. `..module`. + * + * @param mpy A compiled .mpy file (ABI version 6, bytecode only). + * @returns The names of the imported modules. + * @throws {MpyFormatError} If the file is not a .mpy file this code can read. + */ +export function findImportedModules(mpy: Uint8Array): ReadonlySet { + const modules = new Set(); + let pos = 0; + + function readByte(): number { + if (pos >= mpy.length) { + throw new MpyFormatError('unexpected end of .mpy file'); + } + + return mpy[pos++]; + } + + /** Reads a MicroPython variable length unsigned integer. */ + function readUint(): number { + let value = 0; + + for (;;) { + const b = readByte(); + value = (value << 7) | (b & 0x7f); + + if (!(b & 0x80)) { + return value; + } + } + } + + function take(size: number): Uint8Array { + if (size < 0 || pos + size > mpy.length) { + throw new MpyFormatError('unexpected end of .mpy file'); + } + + const bytes = mpy.subarray(pos, pos + size); + pos += size; + + return bytes; + } + + // Header is 'M', ABI version, feature flags, small int bits. + if (mpy.length < 4 || mpy[0] !== 'M'.charCodeAt(0)) { + throw new MpyFormatError('not a .mpy file'); + } + + if (mpy[1] !== mpyAbiVersion) { + throw new MpyFormatError(`unsupported .mpy ABI version: ${mpy[1]}`); + } + + // The top bits of the feature flags are MPY_FEATURE_ENCODE_ARCH(). We only know how + // to walk bytecode, so anything other than MP_NATIVE_ARCH_NONE is not supported. + if (mpy[2] >> 2 !== 0) { + throw new MpyFormatError('.mpy file contains native code'); + } + + pos = 4; + + const nQstr = readUint(); + const nObj = readUint(); + + // The qstr table holds all of the names used by the module. Qstr operands in the + // bytecode are indices into this table. + const decoder = new TextDecoder(); + const qstrTable = new Array(nQstr); + + for (let i = 0; i < nQstr; i++) { + const encodedLength = readUint(); + + if (encodedLength & 1) { + // reference to one of MicroPython's static qstrs + const index = encodedLength >> 1; + const qstr = index < staticQstrs.length ? staticQstrs[index] : undefined; + + if (qstr === undefined || qstr === null) { + throw new MpyFormatError(`bad static qstr index: ${index}`); + } + + qstrTable[i] = qstr; + } else { + const bytes = take(encodedLength >> 1); + take(1); // null terminator + qstrTable[i] = decoder.decode(bytes); + } + } + + // The object table holds constants. We don't need any of them, but they have to be + // stepped over to find the bytecode that follows. + function skipObj(): void { + const type = readByte() as PersistentObj; + + switch (type) { + case PersistentObj.FunTable: + case PersistentObj.None: + case PersistentObj.False: + case PersistentObj.True: + case PersistentObj.Ellipsis: + break; + case PersistentObj.Tuple: { + const n = readUint(); + + for (let i = 0; i < n; i++) { + skipObj(); + } + + break; + } + case PersistentObj.Str: + case PersistentObj.Bytes: + take(readUint()); + take(1); // null terminator + break; + case PersistentObj.Int: + case PersistentObj.Float: + case PersistentObj.Complex: + take(readUint()); + break; + default: + throw new MpyFormatError(`bad .mpy object type: ${type}`); + } + } + + for (let i = 0; i < nObj; i++) { + skipObj(); + } + + /** + * Scans one function's bytecode for import statements. + * + * @param bytecode The function data, starting with the prelude. + */ + function scanBytecode(bytecode: Uint8Array): void { + let ip = 0; + + function preludeByte(): number { + if (ip >= bytecode.length) { + throw new MpyFormatError('unexpected end of .mpy bytecode'); + } + + return bytecode[ip++]; + } + + // Skip the prelude signature (see MP_BC_PRELUDE_SIG_DECODE in py/bc.h). + while (preludeByte() & 0x80) { + // all we need is the length + } + + // Read the prelude size to find where the opcodes start (see + // MP_BC_PRELUDE_SIZE_DECODE in py/bc.h). nInfo covers the source info and + // argument names, nCell covers the closure info. + let nInfo = 0; + let nCell = 0; + + for (let n = 0; ; n++) { + const z = preludeByte(); + nInfo |= ((z & 0x7e) >> 1) << (6 * n); + nCell |= (z & 1) << n; + + if (!(z & 0x80)) { + break; + } + } + + ip += nInfo + nCell; + + // The compiler always pushes the relative import level as a small int + // immediately before MP_BC_IMPORT_NAME, so the most recently loaded small int + // is the number of leading dots. See compile_dotted_as_name() and + // compile_import_from() in py/compile.c. + let importLevel = 0; + + while (ip < bytecode.length) { + const opcode = bytecode[ip]; + const format = (mpBcFormatTable >> (2 * (opcode >> 4))) & 3; + let next = ip + 1; + let arg = 0; + + if (format === mpBcFormatQstr || format === mpBcFormatVarUint) { + arg = bytecode[next] & 0x7f; + + if (opcode === mpBcLoadConstSmallInt && arg & 0x40) { + // sign extend + arg -= 0x80; + } + + while (bytecode[next] & 0x80) { + next++; + arg = (arg << 7) | (bytecode[next] & 0x7f); + } + + next++; + } else if (format === mpBcFormatOffset) { + // we only need the size, not the offset itself + next += bytecode[next] & 0x80 ? 2 : 1; + } + + if ((opcode & mpBcMaskExtraByte) === 0) { + next++; + } + + if (opcode === mpBcLoadConstSmallInt) { + importLevel = arg; + } else if ( + opcode >= mpBcLoadConstSmallIntMulti && + opcode < mpBcLoadConstSmallIntMulti + mpBcLoadConstSmallIntMultiNum + ) { + importLevel = + opcode - + mpBcLoadConstSmallIntMulti - + mpBcLoadConstSmallIntMultiExcess; + } else if (opcode === mpBcImportName) { + const name = arg < qstrTable.length ? qstrTable[arg] : undefined; + + if (name === undefined) { + throw new MpyFormatError(`bad qstr index: ${arg}`); + } + + modules.add('.'.repeat(Math.max(importLevel, 0)) + name); + } + + if (next <= ip) { + // this would mean we misread an opcode size and would loop forever + throw new MpyFormatError(`bad .mpy opcode: ${opcode}`); + } + + ip = next; + } + } + + /** + * Walks one entry of the raw code tree, recursing into nested functions and classes + * so that imports in any scope are found. + */ + function walkRawCode(): void { + const kindLen = readUint(); + const kind = kindLen & 3; + const hasChildren = !!(kindLen & 4); + const funDataLen = kindLen >> 3; + + // 0 is MP_CODE_BYTECODE. Native code is rejected by the header check above, so + // this should be unreachable. + if (kind !== 0) { + throw new MpyFormatError('.mpy file contains native code'); + } + + scanBytecode(take(funDataLen)); + + if (hasChildren) { + const nChildren = readUint(); + + for (let i = 0; i < nChildren; i++) { + walkRawCode(); + } + } + } + + walkRawCode(); + + return modules; +} diff --git a/src/mpy/sagas.test.ts b/src/mpy/sagas.test.ts index 78093b295..c15e594d3 100644 --- a/src/mpy/sagas.test.ts +++ b/src/mpy/sagas.test.ts @@ -1,38 +1,26 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2020-2023 The Pybricks Authors -import path from 'path'; -import { AsyncSaga } from '../../test'; -import { compile, didCompile, didFailToCompile } from './actions'; +import 'core-js/stable/structured-clone'; +import 'fake-indexeddb/auto'; +import 'dexie-observable'; +import { AsyncSaga, uuid } from '../../test'; +import { editorGetValueRequest, editorGetValueResponse } from '../editor/actions'; +import { FileStorageDb } from '../fileStorage'; +import { createCountFunc } from '../utils/iter'; +import { + compile, + didCompile, + didFailToCompile, + mpyCompileMulti6, + mpyDidCompileMulti6, + mpyDidFailToCompileMulti6, +} from './actions'; import mpy from './sagas'; - -const mpyCrossV5Wasm = require.resolve('@pybricks/mpy-cross-v5/build/mpy-cross.wasm'); -const mpyCrossV6Wasm = require.resolve( - '@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm', -); +import { mockMpyCrossWasmPath } from './test-utils'; beforeEach(() => { - // HACK: work around Emscripten + Webpack bugs - // Since we are using jsdom, emscripten thinks we are in a browser and - // sees that the path starts with file:// but just passes this to - // path.normalize() which treats file: as a windows-style drive prefix. - // Also, the webpack import.meta.url doesn't work correctly in the test - // environment either and returns a path relative to the directory where - // it was called rather than the node_modules/ directory. So we have to - // fake the normalization to get the correct path. - jest.spyOn(path, 'normalize').mockImplementation((p) => { - // NB: we can't call require.resolve() here because it would recursively - // call this function via path.normalize()! - if (p.endsWith('@pybricks/mpy-cross-v5/build/mpy-cross.wasm')) { - return mpyCrossV5Wasm; - } - - if (p.endsWith('@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm')) { - return mpyCrossV6Wasm; - } - - return p; - }); + mockMpyCrossWasmPath(); }); afterEach(() => { @@ -71,3 +59,161 @@ test('compiler error works', async () => { await saga.end(); }); + +describe('handleCompileMulti6', () => { + let db: FileStorageDb; + let saga: AsyncSaga; + let nextUuid: () => number; + + beforeEach(async () => { + db = new FileStorageDb('test'); + // the main module is uuid(0), which is what the editor state points at + nextUuid = createCountFunc(); + + saga = new AsyncSaga(mpy, { + fileStorage: db, + nextMessageId: createCountFunc(), + }); + + saga.updateState({ + editor: { isReady: true, activeFileUuid: uuid(0) }, + }); + }); + + afterEach(async () => { + await saga.end(); + db.close(); + + await new Promise((resolve, reject) => { + const request = indexedDB.deleteDatabase('test'); + request.addEventListener('success', resolve); + request.addEventListener('error', reject); + request.addEventListener('blocked', reject); + }); + }); + + /** jsdom's Blob doesn't implement arrayBuffer(), so use FileReader instead. */ + function readBlob(blob: Blob): Promise { + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(reader.result as ArrayBuffer); + reader.onerror = () => reject(reader.error); + reader.readAsArrayBuffer(blob); + }); + } + + /** Adds a module to the simulated user file system. */ + async function addFile(name: string, contents: string): Promise { + await db.metadata.add({ + uuid: uuid(nextUuid()), + path: `${name}.py`, + sha256: '', + viewState: null, + }); + await db._contents.add({ path: `${name}.py`, contents }); + } + + /** + * Runs the multi-mpy6 compile and returns the module names in the resulting + * program, in order. + */ + async function compileMulti6(mainPy: string): Promise { + saga.put(mpyCompileMulti6()); + + // handleCompileMulti6() reads the main file from the editor, not the database + const request = await saga.take(); + expect(editorGetValueRequest.matches(request)).toBeTruthy(); + saga.put( + editorGetValueResponse( + (request as ReturnType).id, + mainPy, + ), + ); + + const action = await saga.take(); + + if (mpyDidFailToCompileMulti6.matches(action)) { + throw new Error(action.error.join('\n')); + } + + expect(mpyDidCompileMulti6.matches(action)).toBeTruthy(); + + // each module is encoded as a uint32 size, a zero-terminated name, then the + // .mpy binary + const blob = (action as ReturnType).file; + const data = new DataView(await readBlob(blob)); + const names: string[] = []; + + for (let offset = 0; offset < data.byteLength; ) { + const size = data.getUint32(offset, true); + offset += 4; + + let end = offset; + while (data.getUint8(end) !== 0) { + end++; + } + + names.push( + new TextDecoder().decode( + new Uint8Array(data.buffer, offset, end - offset), + ), + ); + + offset = end + 1 + size; + } + + return names; + } + + test('a program with no imports contains only the main module', async () => { + await addFile('main', ''); + + await expect(compileMulti6('print("hello!")')).resolves.toEqual(['main']); + }); + + test('imported modules are included', async () => { + await addFile('main', ''); + await addFile('my_lib', 'VALUE = 1\n'); + + await expect( + compileMulti6('import my_lib\nprint(my_lib.VALUE)'), + ).resolves.toEqual(['main', 'my_lib']); + }); + + test('imports are resolved transitively', async () => { + await addFile('main', ''); + await addFile('first', 'import second\n'); + await addFile('second', 'import third\n'); + await addFile('third', 'VALUE = 3\n'); + + await expect(compileMulti6('import first')).resolves.toEqual([ + 'main', + 'first', + 'second', + 'third', + ]); + }); + + test('modules that are not in the file system are assumed to be built in', async () => { + await addFile('main', ''); + + await expect( + compileMulti6('from pybricks.hubs import PrimeHub'), + ).resolves.toEqual(['main']); + }); + + test('circular imports terminate', async () => { + await addFile('main', ''); + await addFile('a', 'import b\n'); + await addFile('b', 'import a\n'); + + await expect(compileMulti6('import a')).resolves.toEqual(['main', 'a', 'b']); + }); + + test('a syntax error in an imported module fails the compile', async () => { + await addFile('main', ''); + await addFile('broken', 'syntax error!\n'); + + await expect(compileMulti6('import broken')).rejects.toThrow(/SyntaxError/); + }); +}); diff --git a/src/mpy/sagas.ts b/src/mpy/sagas.ts index b8c284a1e..dfca189dc 100644 --- a/src/mpy/sagas.ts +++ b/src/mpy/sagas.ts @@ -6,7 +6,7 @@ import { compile as mpyCrossCompileV6 } from '@pybricks/mpy-cross-v6'; import { call, getContext, put, select, takeEvery } from 'typed-redux-saga/macro'; import { editorGetValue } from '../editor/sagaLib'; import { FileContents, FileStorageDb } from '../fileStorage'; -import { findImportedModules, resolveModule } from '../pybricksMicropython/lib'; +import { resolveModule } from '../pybricksMicropython/lib'; import { RootState } from '../reducers'; import { compile, @@ -16,6 +16,7 @@ import { mpyDidCompileMulti6, mpyDidFailToCompileMulti6, } from './actions'; +import { MpyFormatError, findImportedModules } from './mpyImports'; const encoder = new TextEncoder(); @@ -138,76 +139,88 @@ function* handleCompileMulti6(): Generator { ? '__main__' : mainPyPath.replace(/\.[^.]+$/, ''); - const pyFiles = new Map([ - [mainPyName, { path: mainPyPath, contents: mainPyContents }], - ]); + // NB: the URL has to be created outside of the loop or webpack won't be able to + // resolve it as an asset. + const wasmUrl = new URL( + '@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm', + import.meta.url, + ).toString(); + + // Compile the main module, then read the modules it imports back out of the + // compiled bytecode and do the same for each of those, until nothing new is found. + // Each module is compiled exactly once and the order is preserved so that the main + // module comes first in the downloaded program. + const compiled = new Map(); const checkedModules = new Set([mainPyName]); - const uncheckedScripts = new Array(mainPyContents); + const uncompiled = new Array<[string, FileContents]>([ + mainPyName, + { path: mainPyPath, contents: mainPyContents }, + ]); for (;;) { - // parse all unchecked scripts to find imported modules that haven't - // been checked yet + const next = uncompiled.shift(); - const uncheckedModules = new Set(); + if (!next) { + break; + } - for (const uncheckedScript of uncheckedScripts) { - const importedModules = findImportedModules(uncheckedScript); + const [module, py] = next; - for (const m of importedModules) { - if (!checkedModules.has(m)) { - uncheckedModules.add(m); - } - } - } + const result = yield* call(() => + mpyCrossCompileV6(py.path, py.contents, undefined, wasmUrl), + ); - // all of the scripts have been checked now, so clear the unchecked list - uncheckedScripts.length = 0; + if (result.status !== 0 || !result.mpy) { + yield* put(mpyDidFailToCompileMulti6(result.err)); + return; + } - // when no more new modules are found, we are done - if (uncheckedModules.size === 0) { - break; + compiled.set(module, result.mpy); + + let importedModules: ReadonlySet; + + try { + importedModules = findImportedModules(result.mpy); + } catch (err) { + // This means mpy-cross is producing a file format we don't know how to + // read, which would only happen if the mpy-cross dependency changed. Fail + // loudly rather than silently downloading a program with missing modules. + // TODO: error needs to be translated + yield* put( + mpyDidFailToCompileMulti6([ + err instanceof MpyFormatError + ? `failed to read imports of '${py.path}': ${err.message}` + : String(err), + ]), + ); + return; } - // try to resolve unchecked modules in the file system - for (const m of uncheckedModules) { + // try to resolve newly found modules in the file system + for (const m of importedModules) { + if (checkedModules.has(m)) { + continue; + } + + checkedModules.add(m); + const file = yield* call(() => resolveModule(db, m)); - // if found, queue the module to be compiled and to be parsed - // for additional imports + // if not found, the module is assumed to be built in to the firmware if (file) { - pyFiles.set(m, file); - uncheckedScripts.push(file.contents); + uncompiled.push([m, file]); } - - checkedModules.add(m); } } const blobParts: BlobPart[] = []; - for (const [m, py] of pyFiles) { - const result = yield* call(() => - mpyCrossCompileV6( - py.path, - py.contents, - undefined, - new URL( - '@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm', - import.meta.url, - ).toString(), - ), - ); - - if (result.status !== 0 || !result.mpy) { - yield* put(mpyDidFailToCompileMulti6(result.err)); - return; - } - + for (const [module, mpy] of compiled) { // 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); + blobParts.push(encodeUInt32LE(mpy.length)); + blobParts.push(cString(module)); + blobParts.push(mpy); } yield* put(mpyDidCompileMulti6(new Blob(blobParts))); diff --git a/src/mpy/staticQstrs.ts b/src/mpy/staticQstrs.ts new file mode 100644 index 000000000..1ce570808 --- /dev/null +++ b/src/mpy/staticQstrs.ts @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 The Pybricks Authors + +/** + * MicroPython's table of static qstrs. + * + * Qstrs at or below `QSTR_LAST_STATIC` (`MP_QSTR_zip`) are stored in a .mpy file as a + * bare index into this table rather than as a string, so it is needed to recover the + * names of imported modules. Notably `main` is a static qstr, and `main.py` is a very + * common user program name. + * + * Index 0 is `MP_QSTRnull` and is never referenced. + * + * This is `static_qstr_list` from MicroPython's `py/makeqstrdata.py`, with the leading + * null added the same way `tools/mpy-tool.py` does it. The list is frozen by the .mpy + * file format - changing it would invalidate every existing .mpy file - and is verified + * identical between MicroPython v1.19.1 (which `@pybricks/mpy-cross-v6` is built from) + * and pybricks-micropython master. + */ +export const staticQstrs: ReadonlyArray = [ + null, + '', + '__dir__', + '\n', + ' ', + '*', + '/', + '', + '_', + '__call__', + '__class__', + '__delitem__', + '__enter__', + '__exit__', + '__getattr__', + '__getitem__', + '__hash__', + '__init__', + '__int__', + '__iter__', + '__len__', + '__main__', + '__module__', + '__name__', + '__new__', + '__next__', + '__qualname__', + '__repr__', + '__setitem__', + '__str__', + 'ArithmeticError', + 'AssertionError', + 'AttributeError', + 'BaseException', + 'EOFError', + 'Ellipsis', + 'Exception', + 'GeneratorExit', + 'ImportError', + 'IndentationError', + 'IndexError', + 'KeyError', + 'KeyboardInterrupt', + 'LookupError', + 'MemoryError', + 'NameError', + 'NoneType', + 'NotImplementedError', + 'OSError', + 'OverflowError', + 'RuntimeError', + 'StopIteration', + 'SyntaxError', + 'SystemExit', + 'TypeError', + 'ValueError', + 'ZeroDivisionError', + 'abs', + 'all', + 'any', + 'append', + 'args', + 'bool', + 'builtins', + 'bytearray', + 'bytecode', + 'bytes', + 'callable', + 'chr', + 'classmethod', + 'clear', + 'close', + 'const', + 'copy', + 'count', + 'dict', + 'dir', + 'divmod', + 'end', + 'endswith', + 'eval', + 'exec', + 'extend', + 'find', + 'format', + 'from_bytes', + 'get', + 'getattr', + 'globals', + 'hasattr', + 'hash', + 'id', + 'index', + 'insert', + 'int', + 'isalpha', + 'isdigit', + 'isinstance', + 'islower', + 'isspace', + 'issubclass', + 'isupper', + 'items', + 'iter', + 'join', + 'key', + 'keys', + 'len', + 'list', + 'little', + 'locals', + 'lower', + 'lstrip', + 'main', + 'map', + 'micropython', + 'next', + 'object', + 'open', + 'ord', + 'pop', + 'popitem', + 'pow', + 'print', + 'range', + 'read', + 'readinto', + 'readline', + 'remove', + 'replace', + 'repr', + 'reverse', + 'rfind', + 'rindex', + 'round', + 'rsplit', + 'rstrip', + 'self', + 'send', + 'sep', + 'set', + 'setattr', + 'setdefault', + 'sort', + 'sorted', + 'split', + 'start', + 'startswith', + 'staticmethod', + 'step', + 'stop', + 'str', + 'strip', + 'sum', + 'super', + 'throw', + 'to_bytes', + 'tuple', + 'type', + 'update', + 'upper', + 'utf-8', + 'value', + 'values', + 'write', + 'zip', +]; diff --git a/src/mpy/test-utils.ts b/src/mpy/test-utils.ts new file mode 100644 index 000000000..db85a2b9c --- /dev/null +++ b/src/mpy/test-utils.ts @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2020-2026 The Pybricks Authors + +import path from 'path'; + +const mpyCrossV5Wasm = require.resolve('@pybricks/mpy-cross-v5/build/mpy-cross.wasm'); +const mpyCrossV6Wasm = require.resolve( + '@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm', +); + +/** + * Makes the mpy-cross wasm binaries loadable from tests. + * + * HACK: work around Emscripten + Webpack bugs + * Since we are using jsdom, emscripten thinks we are in a browser and + * sees that the path starts with file:// but just passes this to + * path.normalize() which treats file: as a windows-style drive prefix. + * Also, the webpack import.meta.url doesn't work correctly in the test + * environment either and returns a path relative to the directory where + * it was called rather than the node_modules/ directory. So we have to + * fake the normalization to get the correct path. + * + * Call this from `beforeEach()`. + */ +export function mockMpyCrossWasmPath(): void { + jest.spyOn(path, 'normalize').mockImplementation((p) => { + // NB: we can't call require.resolve() here because it would recursively + // call this function via path.normalize()! + if (p.endsWith('@pybricks/mpy-cross-v5/build/mpy-cross.wasm')) { + return mpyCrossV5Wasm; + } + + if (p.endsWith('@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm')) { + return mpyCrossV6Wasm; + } + + return p; + }); +} + +/** + * The wasm path to pass to `mpyCrossCompileV6()` from tests. + * + * This mimics what `new URL('@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm', + * import.meta.url)` produces in the app: a file: URL, so that emscripten reads the file + * instead of trying to fetch it. `mockMpyCrossWasmPath()` rewrites it to the real path. + */ +export const mpyCrossV6WasmPath = + 'file:///@pybricks/mpy-cross-v6/build/mpy-cross-v6.wasm'; diff --git a/src/pybricksMicropython/lib.test.ts b/src/pybricksMicropython/lib.test.ts index db7ae55e4..f480d095f 100644 --- a/src/pybricksMicropython/lib.test.ts +++ b/src/pybricksMicropython/lib.test.ts @@ -3,7 +3,6 @@ import { FileNameValidationResult, - findImportedModules, pythonFileExtension, pythonFileExtensionRegex, validateFileName, @@ -77,72 +76,3 @@ 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 - -# import q -# from q import q -""" -import q -from q import q -""" -''' -import q -from q import q -''' -`; - - 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', - ]), - ); -}); - -test('https://github.com/pybricks/support/issues/873 regression', () => { - const script = ` -from my_module import data - -async def hello(): - print("hello") - -print(data) -`; - - const modules = findImportedModules(script); - - expect(modules).toEqual(new Set(['my_module'])); -}); diff --git a/src/pybricksMicropython/lib.ts b/src/pybricksMicropython/lib.ts index 156bce372..cbba63a23 100644 --- a/src/pybricksMicropython/lib.ts +++ b/src/pybricksMicropython/lib.ts @@ -1,7 +1,6 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2022 The Pybricks Authors -import { parse, walk } from '@pybricks/python-program-analysis'; import type { FileContents, FileStorageDb } from '../fileStorage'; /** The Python file extension ('.py') */ @@ -71,44 +70,6 @@ export function validateFileName( return FileNameValidationResult.IsOk; } -/** - * Finds modules imported by a Python script. - * - * Returns an empty list if there are syntax errors. - * - * @param py A Python Script. - * @returns A list of the names of modules imported by this file. - */ -export function findImportedModules(py: string): ReadonlySet { - const modules = new Set(); - - try { - const tree = parse(py); - - // find all import statements in the syntax tree and collect imported modules - walk(tree, { - onEnterNode(node, _ancestors) { - if (node.type === 'import') { - for (const name of node.names) { - modules.add(name.path); - } - } else if (node.type === 'from') { - modules.add(node.base); - } - }, - }); - } catch (err) { - // istanbul ignore if - if (process.env.NODE_ENV === 'test') { - console.error(err); - } - - // files with syntax errors are ignored - } - - return modules; -} - export async function resolveModule( db: FileStorageDb, module: string, diff --git a/yarn.lock b/yarn.lock index 4f49ce27c..6630d0669 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2980,7 +2980,6 @@ __metadata: "@pybricks/jedi": 1.17.0 "@pybricks/mpy-cross-v5": ^2.0.0 "@pybricks/mpy-cross-v6": ^2.0.0 - "@pybricks/python-program-analysis": ^2.0.0 "@pyodide/webpack-plugin": ^1.3.2 "@reduxjs/toolkit": ^1.9.7 "@shopify/react-i18n": ^7.13.1 @@ -3106,15 +3105,6 @@ __metadata: languageName: unknown linkType: soft -"@pybricks/python-program-analysis@npm:^2.0.0": - version: 2.0.0 - resolution: "@pybricks/python-program-analysis@npm:2.0.0" - dependencies: - lodash: ^4.17.15 - checksum: 2c3b96889b9710e58c46704957d1f935ac56d4a0af081cdf65761fed5bb087292b564cb99ca676df324d61c6306acacceacca0d5605f7e35f0a0f151b6b5b4d0 - languageName: node - linkType: hard - "@pyodide/webpack-plugin@npm:^1.3.2": version: 1.3.2 resolution: "@pyodide/webpack-plugin@npm:1.3.2"