.ipynb support

crates/pyrefly_python/Cargo.toml

@@ -13,7 +13,7 @@ anyhow = "1.0.98"
 dupe = "0.9.1"
 equivalent = "1.0"
 itertools = "0.14.0"
-lsp-types = "0.94.1"
+lsp-types = "0.97.0"
 parse-display = "0.8.2"
 pathdiff = "0.2"
 pyrefly_util = { path = "../pyrefly_util" }
@@ -22,10 +22,10 @@ ruff_python_ast = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a
 ruff_python_parser = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
 ruff_text_size = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
+serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 starlark_map = "0.13.0"
 static_interner = "0.1.1"
 thiserror = "2.0.12"
 
 [dev-dependencies]
-serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 toml = { version = "0.9.2", features = ["preserve_order"] }

crates/pyrefly_python/src/lib.rs

@@ -28,13 +28,14 @@ pub mod module;
 pub mod module_name;
 pub mod module_path;
 pub mod nesting_context;
+pub mod notebook;
 pub mod qname;
 pub mod short_identifier;
 pub mod symbol_kind;
 pub mod sys_info;
 
 /// Suffixes of python files that we can be processed.
-pub const PYTHON_EXTENSIONS: &[&str] = &["py", "pyi"];
+pub const PYTHON_EXTENSIONS: &[&str] = &["py", "pyi", "ipynb"];
 
 /// Suffixes of compiled python modules
 pub const COMPILED_FILE_SUFFIXES: &[&str] = &["pyc", "pyx", "pyd"];

crates/pyrefly_python/src/module_name.rs

@@ -204,7 +204,7 @@ impl ModuleName {
             None => {}
             Some(file_name) => {
                 let splits: Vec<&str> = file_name.rsplitn(2, '.').collect();
-                if splits.len() != 2 || !(splits[0] == "py" || splits[0] == "pyi") {
+                if splits.len() != 2 || !(splits[0] == "py" || splits[0] == "pyi" || splits[0] == "ipynb") {
                     return Err(anyhow::anyhow!(PathConversionError::InvalidExtension {
                         file_name: file_name.to_owned(),
                     }));
@@ -407,8 +407,10 @@ mod tests {
         }
         assert_module_name("foo.py", "foo");
         assert_module_name("foo.pyi", "foo");
+        assert_module_name("foo.ipynb", "foo");
         assert_module_name("foo/bar.py", "foo.bar");
         assert_module_name("foo/bar.pyi", "foo.bar");
+        assert_module_name("foo/bar.ipynb", "foo.bar");
         assert_module_name("foo/bar/__init__.py", "foo.bar");
         assert_module_name("foo/bar/__init__.pyi", "foo.bar");
 

crates/pyrefly_python/src/module_path.rs

@@ -64,6 +64,7 @@ impl ModuleStyle {
         if path.extension() == Some("pyi".as_ref()) {
             ModuleStyle::Interface
         } else {
+            // Both .py and .ipynb are executable
             ModuleStyle::Executable
         }
     }

crates/pyrefly_python/src/notebook.rs

@@ -0,0 +1,134 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Jupyter notebook parsing support.
+//
+// This module extracts Python code from `.ipynb` files (Jupyter notebooks)
+// and converts them into a single Python source file that can be analyzed
+// by the type checker.
+
+use anyhow::Context;
+use anyhow::Result;
+use serde::Deserialize;
+
+/// Represents a Jupyter notebook cell
+#[derive(Debug, Deserialize)]
+struct NotebookCell {
+    cell_type: String,
+    source: NotebookSource,
+}
+
+/// Source can be either a string or an array of strings
+#[derive(Debug, Deserialize)]
+#[serde(untagged)]
+enum NotebookSource {
+    String(String),
+    Array(Vec<String>),
+}
+
+/// Minimal representation of a Jupyter notebook
+#[derive(Debug, Deserialize)]
+struct Notebook {
+    cells: Vec<NotebookCell>,
+}
+
+/// Extracts Python code from a Jupyter notebook JSON string.
+///
+/// This function:
+/// - Parses the notebook JSON
+/// - Extracts only code cells (ignoring markdown cells)
+/// - Concatenates all code into a single Python source string
+/// - Adds cell markers as comments for debugging
+///
+/// # Arguments
+/// * `content` - The raw JSON content of a `.ipynb` file
+///
+/// # Returns
+/// A single Python source string containing all code cells, or an error if parsing fails
+pub fn extract_python_from_notebook(content: &str) -> Result<String> {
+    let notebook: Notebook =
+        serde_json::from_str(content).context("Failed to parse notebook JSON")?;
+
+    let mut python_code = String::new();
+    let mut code_cell_count = 0;
+
+    for cell in notebook.cells.iter() {
+        if cell.cell_type == "code" {
+            code_cell_count += 1;
+            // Add a comment marker for each cell
+            python_code.push_str(&format!("# Cell {}\n", code_cell_count));
+
+            // Extract the source code
+            match &cell.source {
+                NotebookSource::String(s) => {
+                    python_code.push_str(s.as_str());
+                }
+                NotebookSource::Array(lines) => {
+                    for line in lines {
+                        python_code.push_str(line.as_str());
+                    }
+                }
+            }
+
+            // Add spacing between cells
+            if !python_code.ends_with('\n') {
+                python_code.push('\n');
+            }
+            python_code.push('\n');
+        }
+    }
+
+    Ok(python_code)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_python_from_notebook() {
+        let notebook_json = r##"{
+            "cells": [
+                {
+                    "cell_type": "code",
+                    "source": ["def hello():\n", "    return 'world'\n"]
+                },
+                {
+                    "cell_type": "markdown",
+                    "source": ["This is a markdown cell"]
+                },
+                {
+                    "cell_type": "code",
+                    "source": "x = 5"
+                }
+            ]
+        }"##;
+
+        let result = extract_python_from_notebook(notebook_json).unwrap();
+
+        assert!(result.contains("# Cell 1"));
+        assert!(result.contains("def hello():"));
+        assert!(result.contains("return 'world'"));
+        assert!(!result.contains("This is a markdown cell"));
+        assert!(result.contains("# Cell 2"));
+        assert!(result.contains("x = 5"));
+    }
+
+    #[test]
+    fn test_extract_python_from_empty_notebook() {
+        let notebook_json = r#"{"cells": []}"#;
+        let result = extract_python_from_notebook(notebook_json).unwrap();
+        assert_eq!(result, "");
+    }
+
+    #[test]
+    fn test_extract_python_with_invalid_json() {
+        let invalid_json = "not valid json";
+        let result = extract_python_from_notebook(invalid_json);
+        assert!(result.is_err());
+    }
+}

crates/pyrefly_util/Cargo.toml

@@ -23,7 +23,7 @@ human_bytes = "0.4.3"
 ignore = "0.4"
 itertools = "0.14.0"
 lock_free_hashtable = "0.1.1"
-lsp-types = "0.94.1"
+lsp-types = "0.97.0"
 memory-stats = "1.2.0"
 notify = "5"
 parse-display = "0.8.2"

crates/tsp_types/Cargo.toml

@@ -10,6 +10,6 @@ license = "MIT"
 
 [dependencies]
 lsp-server = "0.7.2"
-lsp-types = "0.94.1"
+lsp-types = "0.97.0"
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }

lsp/package.json

@@ -33,7 +33,8 @@
     },
     "main": "./dist/extension",
     "activationEvents": [
-        "onLanguage:python"
+        "onLanguage:python",
+        "onNotebook:jupyter-notebook"
     ],
     "contributes": {
         "languages": [

lsp/src/extension.ts

@@ -41,7 +41,8 @@ async function updateStatusBar() {
   const document = vscode.window.activeTextEditor?.document;
   if (
     document == null ||
-    document.uri.scheme !== 'file' ||
+    (document.uri.scheme !== 'file' &&
+      document.uri.scheme !== 'vscode-notebook-cell') ||
     document.languageId !== 'python'
   ) {
     statusBarItem?.hide();
@@ -180,8 +181,21 @@ export async function activate(context: ExtensionContext) {
   // Options to control the language client
   let clientOptions: LanguageClientOptions = {
     initializationOptions: rawInitialisationOptions,
-    // Register the server for Starlark documents
-    documentSelector: [{scheme: 'file', language: 'python'}],
+    // Register the server for Python documents
+    documentSelector: [
+      {scheme: 'file', language: 'python'},
+      // Support for notebook cells
+      {scheme: 'vscode-notebook-cell', language: 'python'},
+    ],
+    // Support for notebooks
+    notebookDocumentSync: {
+      notebookSelector: [
+        {
+          notebook: {notebookType: 'jupyter-notebook'},
+          cells: [{language: 'python'}],
+        },
+      ],
+    },
     outputChannel: outputChannel,
     middleware: {
       workspace: {

pyrefly/Cargo.toml

@@ -28,7 +28,7 @@ fuzzy-matcher = "0.3.7"
 indicatif = { version = "0.17.6", features = ["futures", "improved_unicode", "rayon", "tokio"] }
 itertools = "0.14.0"
 lsp-server = "0.7.2"
-lsp-types = "0.94.1"
+lsp-types = "0.97.0"
 num-traits = { version = "0.2.19", default-features = false }
 parse-display = "0.8.2"
 paste = "1.0.14"