basic ipynb support for type checking

Summary:
This is an absolute minimal implementation, s.t `pyrefly check <some ipynb file>` emits errors.

The error positions are all wrong, because they used the concatenated file. This will be fixed later.

Reviewed By: kinto0

Differential Revision: D85438479

fbshipit-source-id: 99e96b510b708770b16f7c83e7b92b0156f26359

Author: yangdanny97

Cargo.lock

@@ -18,8 +18,10 @@ parse-display = "0.8.2"
 pathdiff = "0.2"
 pyrefly_util = { path = "../pyrefly_util" }
 regex = "1.11.1"
+ruff_notebook = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
 ruff_python_ast = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
 ruff_python_parser = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
+ruff_source_file = { 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"] }
 starlark_map = "0.13.0"

crates/pyrefly_python/Cargo.toml

@@ -33,8 +33,8 @@ 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"];
+/// Suffixes of python files that can be processed.
+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/lib.rs

@@ -11,7 +11,10 @@ use dupe::Dupe;
 use pyrefly_util::arc_id::ArcId;
 use pyrefly_util::lined_buffer::DisplayPos;
 use pyrefly_util::lined_buffer::DisplayRange;
+use pyrefly_util::lined_buffer::LineNumber;
 use pyrefly_util::lined_buffer::LinedBuffer;
+use ruff_notebook::Notebook;
+use ruff_source_file::OneIndexed;
 use ruff_text_size::TextRange;
 use ruff_text_size::TextSize;
 
@@ -32,6 +35,7 @@ struct ModuleInner {
     ignore: Ignore,
     is_generated: bool,
     contents: LinedBuffer,
+    notebook: Option<Box<Notebook>>,
 }
 
 impl Module {
@@ -46,6 +50,22 @@ impl Module {
             ignore,
             is_generated,
             contents,
+            notebook: None,
+        }))
+    }
+
+    pub fn new_notebook(name: ModuleName, path: ModulePath, notebook: Notebook) -> Self {
+        let contents: Arc<String> = Arc::from(notebook.source_code().to_owned());
+        let ignore = Ignore::new(&contents);
+        let is_generated = contents.contains(GENERATED_TOKEN);
+        let contents = LinedBuffer::new(contents);
+        Self(ArcId::new(ModuleInner {
+            name,
+            path,
+            ignore,
+            is_generated,
+            contents,
+            notebook: Some(Box::new(notebook)),
         }))
     }
 
@@ -87,6 +107,10 @@ impl Module {
         &self.0.path
     }
 
+    pub fn is_notebook(&self) -> bool {
+        self.0.notebook.is_some()
+    }
+
     pub fn is_generated(&self) -> bool {
         self.0.is_generated
     }
@@ -116,6 +140,10 @@ impl Module {
     pub fn ignore(&self) -> &Ignore {
         &self.0.ignore
     }
+
+    pub fn notebook(&self) -> Option<&Notebook> {
+        self.0.notebook.as_deref()
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -129,3 +157,21 @@ impl TextRangeWithModule {
         Self { module, range }
     }
 }
+
+/// Given a one-indexed row and column in the concatenated source,
+/// return the cell number, and the row/column in that cell.
+pub fn map_notebook_position(
+    notebook: &Notebook,
+    position: DisplayPos,
+) -> Option<(OneIndexed, DisplayPos)> {
+    let index = notebook.index();
+    let cell = index.cell(position.line.to_one_indexed())?;
+    let cell_row = index
+        .cell_row(position.line.to_one_indexed())
+        .unwrap_or(OneIndexed::MIN);
+    let display_pos = DisplayPos {
+        line: LineNumber::from_one_indexed(cell_row),
+        column: position.column,
+    };
+    Some((cell, display_pos))
+}

crates/pyrefly_python/src/module.rs

@@ -208,7 +208,9 @@ 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(),
                     }));
@@ -417,8 +419,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_name.rs

@@ -95,6 +95,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/module_path.rs

@@ -142,7 +142,7 @@ impl Glob {
     }
 
     fn is_python_extension(ext: Option<&OsStr>) -> bool {
-        ext.is_some_and(|e| e == "py" || e == "pyi" || e == "pyw")
+        ext.is_some_and(|e| e == "py" || e == "pyi" || e == "pyw" || e == "ipynb")
     }
 
     /// Returns true if the given file should be included in results.

crates/pyrefly_util/src/globs.rs

@@ -227,7 +227,11 @@ impl LineNumber {
         self.0.get() - 1
     }
 
-    fn to_one_indexed(self) -> OneIndexed {
+    pub fn from_one_indexed(x: OneIndexed) -> Self {
+        Self(NonZeroU32::new(x.get().try_into().unwrap()).unwrap())
+    }
+
+    pub fn to_one_indexed(self) -> OneIndexed {
         OneIndexed::new(self.0.get() as usize).unwrap()
     }
 

crates/pyrefly_util/src/lined_buffer.rs

@@ -42,6 +42,7 @@ pyrefly_util = { path = "../crates/pyrefly_util" }
 rayon = "1.11.0"
 regex = "1.11.1"
 ruff_annotate_snippets = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
+ruff_notebook = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
 ruff_python_ast = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668", features = ["serde"] }
 ruff_python_parser = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }
 ruff_source_file = { git = "https://github.com/astral-sh/ruff/", rev = "9bee8376a17401f9736b45fdefffb62edc2f1668" }