diff --git a/crates/ruff/src/commands/analyze_graph.rs b/crates/ruff/src/commands/analyze_graph.rs index ffd7cc2d15..d4085e8ed0 100644 --- a/crates/ruff/src/commands/analyze_graph.rs +++ b/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); diff --git a/crates/ruff/tests/analyze_graph.rs b/crates/ruff/tests/analyze_graph.rs index 2c300029ea..993ebf3b59 100644 --- a/crates/ruff/tests/analyze_graph.rs +++ b/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(()) +} diff --git a/crates/ruff_graph/src/lib.rs b/crates/ruff_graph/src/lib.rs index eaf307018d..377f1e89e9 100644 --- a/crates/ruff_graph/src/lib.rs +++ b/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 { - // 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()));