Typing and formatting (GEN-290, GEN-279)

.github/workflows/pyright.yml

@@ -0,0 +1,36 @@
+name: Pyright type checks
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  pyright:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.11.5
+
+      - id: auth
+        uses: google-github-actions/auth@v2
+        with:
+          credentials_json: "${{ secrets.ARTIFACT_REGISTRY_KEY }}"
+
+      - name: Set up Cloud SDK
+        uses: google-github-actions/setup-gcloud@v2
+
+      - name: Install and configure Poetry
+        uses: snok/install-poetry@v1
+
+      - run: poetry self update && poetry self add keyrings.google-artifactregistry-auth
+
+      - run: poetry install
+
+      - uses: jakebailey/pyright-action@v1
+        with:
+          version: 1.1.358

.github/workflows/release.yml

@@ -35,8 +35,8 @@ jobs:
       - name: Install dependencies with Yarn
         run: yarn install --frozen-lockfile --production=false
 
-      - name: Compile the project
-        run: ./compile.sh
+      - name: build the javascript bundle
+        run: yarn build
 
       - name: Install and configure Poetry
         uses: snok/install-poetry@v1

.github/workflows/ruff.yml

@@ -0,0 +1,22 @@
+name: Ruff linter
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  ruff:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v4
+      with:
+        python-version: 3.11.5
+
+    - uses: chartboost/ruff-action@v1
+      with:
+        version: 0.3.0
+        args: --output-format github

.gitignore

@@ -1,4 +1,4 @@
 **/*compiled.js
 *.ipynb
 *.slides.html
-meta.json
+meta.json

.pre-commit-config.yaml

@@ -0,0 +1,42 @@
+ci:
+    skip: [ruff, pyright]
+
+default_install_hook_types:
+  - pre-commit
+  - post-checkout
+  - post-merge
+
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: check-yaml
+        args: [--unsafe]
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.7.1
+    hooks:
+      - id: nbstripout
+
+  - repo: https://github.com/python-poetry/poetry
+    rev: 1.8.0
+    hooks:
+      - id: poetry-check
+      - id: poetry-install
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.4.7
+    hooks:
+      - id: ruff
+        types_or: [ python, pyi, jupyter ]
+        args: [ --fix ]
+
+      - id: ruff-format
+        types_or: [ python, pyi, jupyter ]
+
+  - repo: https://github.com/RobertCraigie/pyright-python
+    rev: v1.1.368
+    hooks:
+    - id: pyright

DEVELOPING.md

@@ -0,0 +1,34 @@
+# Developer's Guide
+
+This guide describes how to complete various tasks you'll encounter when working
+on the GenStudio codebase.
+
+### Commit Hooks
+
+We use [pre-commit](https://pre-commit.com/) to manage a series of git
+pre-commit hooks for the project; for example, each time you commit code, the
+hooks will make sure that your python is formatted properly. If your code isn't,
+the hook will format it, so when you try to commit the second time you'll get
+past the hook.
+
+All hooks are defined in `.pre-commit-config.yaml`. To install these hooks,
+install `pre-commit` if you don't yet have it. I prefer using
+[pipx](https://github.com/pipxproject/pipx) so that `pre-commit` stays globally
+available.
+
+```bash
+pipx install pre-commit
+```
+
+Then install the hooks with this command:
+
+```bash
+pre-commit install
+```
+
+Now they'll run on every commit. If you want to run them manually, run the
+following command:
+
+```bash
+pre-commit run --all-files
+```

README.md

@@ -1,4 +1,4 @@
-# Gen Studio 
+# Gen Studio
 _Visualization tools for GenJAX._
 
 -----
@@ -14,10 +14,10 @@ Key features:
 
 For runnable examples, refer to `notebooks/plot_examples.py`. Detailed examples of all mark types are available at [Observable Plot](https://observablehq.com/plot/).
 
-## Installation 
+## Installation
 
 GenStudio is available in the same artifact registry as GenJAX. Follow [these instructions](https://github.com/probcomp/genjax?tab=readme-ov-file#quickstart) and replace `genjax` with `genstudio` as the package name.
 
 ## Dev
 
-Run `yarn watch` to compile the javascript bundle.
+Run `yarn watch` to compile the javascript bundle.

notebooks/anywidget_jax_callback.js

@@ -4,7 +4,7 @@ import htm from 'https://esm.sh/htm@3.1.1'
 
 const {useState, useEffect} = React
 const html = htm.bind(React.createElement)
-  
+
   const useCustomMessages = (model) => {
       const [messages, setMessages] = useState([])
       const handleMessage = (message, info) => {
@@ -15,7 +15,7 @@ const html = htm.bind(React.createElement)
       };
       useEffect(() => {
         model.on("msg:custom", handleMessage)
-        return () => model.off("msg:custom", handleMessage)    
+        return () => model.off("msg:custom", handleMessage)
       }, [model])
       return messages
   }
@@ -27,12 +27,12 @@ const html = htm.bind(React.createElement)
         useEffect(() => {
             if (messages.length > 0) {
                 // Send a response to python
-                invoke("_receive_message", messages[messages.length - 1], [])    
+                invoke("_receive_message", messages[messages.length - 1], [])
             }
-            }, 
+            },
         [messages.join()])
         return html`<div>
                     Messages: ${messages.join(', ')}
                 </div>`
     }
-  export default {render: createRender(App)}
+  export default {render: createRender(App)}

notebooks/anywidget_jax_callback.py

@@ -1,7 +1,7 @@
 # %%
 # Example of using JAX's io_callback to send messages to an AnyWidget instance.
-# The widget responds, but python won't receive the message until after the 
-# current cell has evaluated. 
+# The widget responds, but python won't receive the message until after the
+# current cell has evaluated.
 
 from jax.experimental import io_callback
 import jax
@@ -10,28 +10,33 @@
 
 # %%
 
+
 class Widget(anywidget.AnyWidget):
     _esm = "anywidget_jax_callback.js"
 
     def __init__(self):
         super().__init__()
 
-    @anywidget.experimental.command
+    @anywidget.experimental.command  # type: ignore
     def _receive_message(self, msg, buffers):
         print(f"Received message from js: {msg}")
         return ["ok", []]
 
+
 w = Widget()
 w
 # %%
 
+
 def effectful_fn(x):
     print(f"Performing an effect with {x}")
-    w.send({'kind': 'genstudio', 'content': x.tolist()})
+    w.send({"kind": "genstudio", "content": x.tolist()})
     return x
 
+
 @jax.jit
 def numpy_random_like(x):
     return io_callback(effectful_fn, x, x)
 
+
 jax.vmap(numpy_random_like)(jnp.zeros(5))

notebooks/anywidget_response_time.py

@@ -1,10 +1,11 @@
 # %%
-# Simple demo of sending a message to an AnyWidget instance and measuring the round trip 
+# Simple demo of sending a message to an AnyWidget instance and measuring the round trip
 # response time.
 
 
 import anywidget
 
+
 class Widget(anywidget.AnyWidget):
     _esm = """
   import { html, render, useState } from 'https://esm.sh/htm/preact/standalone'
@@ -30,7 +31,7 @@ class Widget(anywidget.AnyWidget):
   export default {render: renderWidget}
   """
 
-    @anywidget.experimental.command
+    @anywidget.experimental.command  # type: ignore
     def ping(self, msg, buffers):
         return "pong", None
 

notebooks/intro/compile.sh

@@ -6,4 +6,4 @@ fswatch -o intro_slides.py | while read; do
   echo "Sync and convert complete."
 done
 
-# jupyter nbconvert --to slides --execute --ExecutePreprocessor.kernel_name=studio intro_slides.ipynb
+# jupyter nbconvert --to slides --execute --ExecutePreprocessor.kernel_name=studio intro_slides.ipynb