[ruff]: Make ruff analyze graph work with jupyter notebooks (#21161)

Co-authored-by: Gautham Venkataraman <gautham@dexterenergy.ai>
Co-authored-by: Micha Reiser <micha@reiser.io>

Author: gauthsvenkat

crates/ruff/src/commands/analyze_graph.rs

@@ -7,6 +7,7 @@ use path_absolutize::CWD;
 use ruff_db::system::{SystemPath, SystemPathBuf};
 use ruff_graph::{Direction, ImportMap, ModuleDb, ModuleImports};
 use ruff_linter::package::PackageRoot;
+use ruff_linter::source_kind::SourceKind;
 use ruff_linter::{warn_user, warn_user_once};
 use ruff_python_ast::{PySourceType, SourceType};
 use ruff_workspace::resolver::{ResolvedFile, match_exclusion, python_files_in_path};
@@ -127,10 +128,6 @@ pub(crate) fn analyze_graph(
                     },
                     Some(language) => PySourceType::from(language),
                 };
-                if matches!(source_type, PySourceType::Ipynb) {
-                    debug!("Ignoring Jupyter notebook: {}", path.display());
-                    continue;
-                }
 
                 // Convert to system paths.
                 let Ok(package) = package.map(SystemPathBuf::from_path_buf).transpose() else {
@@ -147,13 +144,34 @@ pub(crate) fn analyze_graph(
                 let root = root.clone();
                 let result = inner_result.clone();
                 scope.spawn(move |_| {
+                    // Extract source code (handles both .py and .ipynb files)
+                    let source_kind = match SourceKind::from_path(path.as_std_path(), source_type) {
+                        Ok(Some(source_kind)) => source_kind,
+                        Ok(None) => {
+                            debug!("Skipping non-Python notebook: {path}");
+                            return;
+                        }
+                        Err(err) => {
+                            warn!("Failed to read source for {path}: {err}");
+                            return;
+                        }
+                    };
+
+                    let source_code = source_kind.source_code();
+
                     // Identify any imports via static analysis.
-                    let mut imports =
-                        ModuleImports::detect(&db, &path, package.as_deref(), string_imports)
-                            .unwrap_or_else(|err| {
-                                warn!("Failed to generate import map for {path}: {err}");
-                                ModuleImports::default()
-                            });
+                    let mut imports = ModuleImports::detect(
+                        &db,
+                        source_code,
+                        source_type,
+                        &path,
+                        package.as_deref(),
+                        string_imports,
+                    )
+                    .unwrap_or_else(|err| {
+                        warn!("Failed to generate import map for {path}: {err}");
+                        ModuleImports::default()
+                    });
 
                     debug!("Discovered {} imports for {}", imports.len(), path);
 

crates/ruff/tests/analyze_graph.rs

@@ -653,3 +653,133 @@ fn venv() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn notebook_basic() -> Result<()> {
+    let tempdir = TempDir::new()?;
+    let root = ChildPath::new(tempdir.path());
+
+    root.child("ruff").child("__init__.py").write_str("")?;
+    root.child("ruff")
+        .child("a.py")
+        .write_str(indoc::indoc! {r#"
+        def helper():
+            pass
+    "#})?;
+
+    // Create a basic notebook with a simple import
+    root.child("notebook.ipynb").write_str(indoc::indoc! {r#"
+        {
+          "cells": [
+            {
+              "cell_type": "code",
+              "execution_count": null,
+              "metadata": {},
+              "outputs": [],
+              "source": [
+                "from ruff.a import helper"
+              ]
+            }
+          ],
+          "metadata": {
+            "language_info": {
+              "name": "python",
+              "version": "3.12.0"
+            }
+          },
+          "nbformat": 4,
+          "nbformat_minor": 5
+        }
+    "#})?;
+
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec(),
+    }, {
+        assert_cmd_snapshot!(command().current_dir(&root), @r###"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        {
+          "notebook.ipynb": [
+            "ruff/a.py"
+          ],
+          "ruff/__init__.py": [],
+          "ruff/a.py": []
+        }
+
+        ----- stderr -----
+        "###);
+    });
+
+    Ok(())
+}
+
+#[test]
+fn notebook_with_magic() -> Result<()> {
+    let tempdir = TempDir::new()?;
+    let root = ChildPath::new(tempdir.path());
+
+    root.child("ruff").child("__init__.py").write_str("")?;
+    root.child("ruff")
+        .child("a.py")
+        .write_str(indoc::indoc! {r#"
+        def helper():
+            pass
+    "#})?;
+
+    // Create a notebook with IPython magic commands and imports
+    root.child("notebook.ipynb").write_str(indoc::indoc! {r#"
+        {
+          "cells": [
+            {
+              "cell_type": "code",
+              "execution_count": null,
+              "metadata": {},
+              "outputs": [],
+              "source": [
+                "%load_ext autoreload\n",
+                "%autoreload 2"
+              ]
+            },
+            {
+              "cell_type": "code",
+              "execution_count": null,
+              "metadata": {},
+              "outputs": [],
+              "source": [
+                "from ruff.a import helper"
+              ]
+            }
+          ],
+          "metadata": {
+            "language_info": {
+              "name": "python",
+              "version": "3.12.0"
+            }
+          },
+          "nbformat": 4,
+          "nbformat_minor": 5
+        }
+    "#})?;
+
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec(),
+    }, {
+        assert_cmd_snapshot!(command().current_dir(&root), @r###"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        {
+          "notebook.ipynb": [
+            "ruff/a.py"
+          ],
+          "ruff/__init__.py": [],
+          "ruff/a.py": []
+        }
+
+        ----- stderr -----
+        "###);
+    });
+
+    Ok(())
+}

crates/ruff_graph/src/lib.rs

@@ -3,8 +3,9 @@ use std::collections::{BTreeMap, BTreeSet};
 use anyhow::Result;
 
 use ruff_db::system::{SystemPath, SystemPathBuf};
+use ruff_python_ast::PySourceType;
 use ruff_python_ast::helpers::to_module_path;
-use ruff_python_parser::{Mode, ParseOptions, parse};
+use ruff_python_parser::{ParseOptions, parse};
 
 use crate::collector::Collector;
 pub use crate::db::ModuleDb;
@@ -24,13 +25,14 @@ impl ModuleImports {
     /// Detect the [`ModuleImports`] for a given Python file.
     pub fn detect(
         db: &ModuleDb,
+        source: &str,
+        source_type: PySourceType,
         path: &SystemPath,
         package: Option<&SystemPath>,
         string_imports: StringImports,
     ) -> Result<Self> {
-        // Read and parse the source code.
-        let source = std::fs::read_to_string(path)?;
-        let parsed = parse(&source, ParseOptions::from(Mode::Module))?;
+        // Parse the source code.
+        let parsed = parse(source, ParseOptions::from(source_type))?;
 
         let module_path =
             package.and_then(|package| to_module_path(package.as_std_path(), path.as_std_path()));