allow installing jiter, and thereby openai and anthropic (#43)

Author: samuelcolvin

README.md

@@ -12,10 +12,50 @@ If you choose to save code, it's stored in CloudFlare's R2 object storage, and s
 
 Dependencies are installed when code is run.
 
-Dependencies can be either:
+Dependencies can be defined in one of two ways:
 
-- defined via [inline script metadata](https://packaging.python.org/en/latest/specifications/inline-script-metadata/#inline-script-metadata) — e.g. a comment at the top of the file, as used by [uv](https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies)
-- or, inferred from imports in the code — e.g. `import pydantic` will install the `pydantic` package
+### Inferred from imports
+
+If there's no metadata, dependencies are inferred from imports in the code.
+
+```py
+import pydantic
+
+class Model(pydantic.BaseModel):
+    x: int
+
+print(Model(x='42'))
+```
+
+### Inline script metadata
+
+As introduced in PEP 723, explained [here](https://packaging.python.org/en/latest/specifications/inline-script-metadata/#inline-script-metadata), and popularised by [uv](https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies) — dependencies can be defined in a comment at the top of the file.
+
+This allows use of dependencies that aren't imported in the code, and is more explicit.
+
+```py
+# /// script
+# dependencies = ["pydantic", "email-validator"]
+# ///
+import pydantic
+
+class Model(pydantic.BaseModel):
+    email: pydantic.EmailStr
+
+print(Model(email='[email protected]'))
+```
+
+It also allows version to be pinned for non-binary packages (Pyodide only supports a single version for the binary packages it supports, like `pydantic` and `numpy`).
+
+```py
+# /// script
+# dependencies = ["rich<13"]
+# ///
+import rich
+from importlib.metadata import version
+
+rich.print(f'[red]Rich version:[/red] [blue]{version('rich')}[/blue]')
+```
 
 ### Sandbox via link
 

src/frontend/src/app.tsx

@@ -24,8 +24,7 @@ export default function () {
       if (data.kind == 'print') {
         newTerminalOutput = true
         for (const chunk of data.data) {
-          const arr = new Uint8Array(chunk)
-          terminalOutput += decoder.decode(arr)
+          terminalOutput += decoder.decode(chunk)
         }
       } else if (data.kind == 'status') {
         setStatus(data.message)

src/frontend/src/install_dependencies.py

@@ -15,8 +15,16 @@
 from pathlib import Path
 from typing import Any, TypedDict, Iterable, Literal
 import importlib.util
+from urllib.parse import urlparse
 
 import tomllib
+from packaging.tags import parse_tag  # noqa
+from packaging.version import Version  # noqa
+
+import micropip  # noqa
+from micropip import transaction  # noqa
+from micropip.wheelinfo import WheelInfo  # noqa
+
 from pyodide.code import find_imports  # noqa
 import pyodide_js  # noqa
 
@@ -41,6 +49,33 @@ class Error:
     kind: Literal['error'] = 'error'
 
 
+# This is a temporary hack to install jiter from a URL until
+# https://github.com/pyodide/pyodide/pull/5388 is released.
+real_find_wheel = transaction.find_wheel
+
+
+def custom_find_wheel(metadata: Any, req: Any) -> Any:
+    if metadata.name == 'jiter':
+        known_version = Version('0.8.2')
+        if known_version in metadata.releases:
+            tag = 'cp312-cp312-emscripten_3_1_58_wasm32'
+            filename = f'{metadata.name}-{known_version}-{tag}.whl'
+            url = f'https://files.pydantic.run/{filename}'
+            return WheelInfo(
+                name=metadata.name,
+                version=known_version,
+                filename=filename,
+                build=(),
+                tags=frozenset({parse_tag(tag)}),
+                url=url,
+                parsed_url=urlparse(url),
+            )
+    return real_find_wheel(metadata, req)
+
+
+transaction.find_wheel = custom_find_wheel
+
+
 async def install_deps(files: list[File]) -> Success | Error:
     sys.setrecursionlimit(400)
     cwd = Path.cwd()
@@ -84,7 +119,6 @@ async def install_deps(files: list[File]) -> Success | Error:
         if install_ssl:
             install_dependencies.append('ssl')
 
-        import micropip  # noqa
         with _micropip_logging() as logs_filename:
             try:
                 await micropip.install(install_dependencies, keep_going=True)

src/frontend/src/types.d.ts

@@ -11,7 +11,7 @@ export interface RunCode {
 
 export interface Print {
   kind: 'print'
-  data: ArrayBuffer[]
+  data: Uint8Array[]
 }
 export interface Message {
   kind: 'status' | 'error' | 'installed'

src/frontend/src/worker.ts

@@ -52,6 +52,7 @@ self.onmessage = async ({ data }: { data: RunCode }) => {
     postPrint()
     post({ kind: 'status', message: `${msg}ran code in ${asMs(execTime)}` })
   } catch (err) {
+    postPrint()
     console.warn(err)
     post({ kind: 'status', message: `${msg}Error occurred` })
     post({ kind: 'error', message: formatError(err) })
@@ -156,12 +157,13 @@ function makeTtyOps() {
   }
 }
 
-let chunks: ArrayBuffer[] = []
+let chunks: Uint8Array[] = []
 let last_post = 0
 
 function print(tty: any) {
-  if (tty.output && tty.output.length > 0) {
-    chunks.push(tty.output)
+  const output: number[] | null = tty.output
+  if (output && output.length > 0) {
+    chunks.push(new Uint8Array(output))
     tty.output = []
     const now = performance.now()
     if (now - last_post > 100) {