[red-knot] Add notebook support (#12338)

crates/ruff_db/src/source.rs

@@ -1,47 +1,83 @@
-use countme::Count;
-use ruff_source_file::LineIndex;
-use salsa::DebugWithDb;
 use std::ops::Deref;
 use std::sync::Arc;
 
+use countme::Count;
+use salsa::DebugWithDb;
+
+use ruff_notebook::Notebook;
+use ruff_python_ast::PySourceType;
+use ruff_source_file::LineIndex;
+
 use crate::files::File;
 use crate::Db;
 
-/// Reads the content of file.
+/// Reads the source text of a python text file (must be valid UTF8) or notebook.
 #[salsa::tracked]
 pub fn source_text(db: &dyn Db, file: File) -> SourceText {
     let _span = tracing::trace_span!("source_text", ?file).entered();
 
-    let content = file.read_to_string(db);
+    if let Some(path) = file.path(db).as_system_path() {
+        if path.extension().is_some_and(|extension| {
+            PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
+        }) {
+            // TODO(micha): Proper error handling and emit a diagnostic. Tackle it together with `source_text`.
+            let notebook = db.system().read_to_notebook(path).unwrap_or_else(|error| {
+                tracing::error!("Failed to load notebook: {error}");
+                Notebook::empty()
+            });
+
+            return SourceText {
+                inner: Arc::new(SourceTextInner {
+                    kind: SourceTextKind::Notebook(notebook),
+                    count: Count::new(),
+                }),
+            };
+        }
+    };
+
+    let content = file.read_to_string(db).unwrap_or_else(|error| {
+        tracing::error!("Failed to load file: {error}");
+        String::default()
+    });
 
     SourceText {
-        inner: Arc::from(content),
-        count: Count::new(),
+        inner: Arc::new(SourceTextInner {
+            kind: SourceTextKind::Text(content),
+            count: Count::new(),
+        }),
     }
 }
 
-/// Computes the [`LineIndex`] for `file`.
-#[salsa::tracked]
-pub fn line_index(db: &dyn Db, file: File) -> LineIndex {
-    let _span = tracing::trace_span!("line_index", file = ?file.debug(db)).entered();
-
-    let source = source_text(db, file);
-
-    LineIndex::from_source_text(&source)
-}
-
-/// The source text of a [`File`].
+/// The source text of a file containing python code.
+///
+/// The file containing the source text can either be a text file or a notebook.
 ///
 /// Cheap cloneable in `O(1)`.
 #[derive(Clone, Eq, PartialEq)]
 pub struct SourceText {
-    inner: Arc<str>,
-    count: Count<Self>,
+    inner: Arc<SourceTextInner>,
 }
 
 impl SourceText {
+    /// Returns the python code as a `str`.
     pub fn as_str(&self) -> &str {
-        &self.inner
+        match &self.inner.kind {
+            SourceTextKind::Text(source) => source,
+            SourceTextKind::Notebook(notebook) => notebook.source_code(),
+        }
+    }
+
+    /// Returns the underlying notebook if this is a notebook file.
+    pub fn as_notebook(&self) -> Option<&Notebook> {
+        match &self.inner.kind {
+            SourceTextKind::Notebook(notebook) => Some(notebook),
+            SourceTextKind::Text(_) => None,
+        }
+    }
+
+    /// Returns `true` if this is a notebook source file.
+    pub fn is_notebook(&self) -> bool {
+        matches!(&self.inner.kind, SourceTextKind::Notebook(_))
     }
 }
 
@@ -55,20 +91,54 @@ impl Deref for SourceText {
 
 impl std::fmt::Debug for SourceText {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_tuple("SourceText").field(&self.inner).finish()
+        let mut dbg = f.debug_tuple("SourceText");
+
+        match &self.inner.kind {
+            SourceTextKind::Text(text) => {
+                dbg.field(text);
+            }
+            SourceTextKind::Notebook(notebook) => {
+                dbg.field(notebook);
+            }
+        }
+
+        dbg.finish()
     }
 }
 
+#[derive(Eq, PartialEq)]
+struct SourceTextInner {
+    count: Count<SourceText>,
+    kind: SourceTextKind,
+}
+
+#[derive(Eq, PartialEq)]
+enum SourceTextKind {
+    Text(String),
+    Notebook(Notebook),
+}
+
+/// Computes the [`LineIndex`] for `file`.
+#[salsa::tracked]
+pub fn line_index(db: &dyn Db, file: File) -> LineIndex {
+    let _span = tracing::trace_span!("line_index", file = ?file.debug(db)).entered();
+
+    let source = source_text(db, file);
+
+    LineIndex::from_source_text(&source)
+}
+
 #[cfg(test)]
 mod tests {
     use salsa::EventKind;
 
+    use ruff_source_file::OneIndexed;
+    use ruff_text_size::TextSize;
+
     use crate::files::system_path_to_file;
     use crate::source::{line_index, source_text};
     use crate::system::{DbWithTestSystem, SystemPath};
     use crate::tests::TestDb;
-    use ruff_source_file::OneIndexed;
-    use ruff_text_size::TextSize;
 
     #[test]
     fn re_runs_query_when_file_revision_changes() -> crate::system::Result<()> {
@@ -79,11 +149,11 @@ mod tests {
 
         let file = system_path_to_file(&db, path).unwrap();
 
-        assert_eq!(&*source_text(&db, file), "x = 10");
+        assert_eq!(source_text(&db, file).as_str(), "x = 10");
 
         db.write_file(path, "x = 20".to_string()).unwrap();
 
-        assert_eq!(&*source_text(&db, file), "x = 20");
+        assert_eq!(source_text(&db, file).as_str(), "x = 20");
 
         Ok(())
     }
@@ -97,13 +167,13 @@ mod tests {
 
         let file = system_path_to_file(&db, path).unwrap();
 
-        assert_eq!(&*source_text(&db, file), "x = 10");
+        assert_eq!(source_text(&db, file).as_str(), "x = 10");
 
         // Change the file permission only
         file.set_permissions(&mut db).to(Some(0o777));
 
         db.clear_salsa_events();
-        assert_eq!(&*source_text(&db, file), "x = 10");
+        assert_eq!(source_text(&db, file).as_str(), "x = 10");
 
         let events = db.take_salsa_events();
 
@@ -123,14 +193,54 @@ mod tests {
 
         let file = system_path_to_file(&db, path).unwrap();
         let index = line_index(&db, file);
-        let text = source_text(&db, file);
+        let source = source_text(&db, file);
 
         assert_eq!(index.line_count(), 2);
         assert_eq!(
-            index.line_start(OneIndexed::from_zero_indexed(0), &text),
+            index.line_start(OneIndexed::from_zero_indexed(0), source.as_str()),
             TextSize::new(0)
         );
 
         Ok(())
     }
+
+    #[test]
+    fn notebook() -> crate::system::Result<()> {
+        let mut db = TestDb::new();
+
+        let path = SystemPath::new("test.ipynb");
+        db.write_file(
+            path,
+            r#"
+{
+    "cells": [{"cell_type": "code", "source": ["x = 10"], "metadata": {}, "outputs": []}],
+    "metadata": {
+        "kernelspec": {
+            "display_name": "Python (ruff)",
+            "language": "python",
+            "name": "ruff"
+        },
+        "language_info": {
+            "file_extension": ".py",
+            "mimetype": "text/x-python",
+            "name": "python",
+            "nbconvert_exporter": "python",
+            "pygments_lexer": "ipython3",
+            "version": "3.11.3"
+        }
+     },
+     "nbformat": 4,
+     "nbformat_minor": 4
+}"#,
+        )?;
+
+        let file = system_path_to_file(&db, path).unwrap();
+        let source = source_text(&db, file);
+
+        assert!(source.is_notebook());
+        assert_eq!(source.as_str(), "x = 10\n");
+        assert!(source.as_notebook().is_some());
+
+        Ok(())
+    }
 }