Implement import extraction for notebooks in Rust

Cargo.toml

@@ -18,6 +18,7 @@ rayon = "1.9.0"
 regex = "1.10.3"
 rustpython-ast = { version = "0.3.0", features = ["visitor"] }
 rustpython-parser = "0.3.0"
+serde_json = "1.0.114"
 
 [profile.release]
 lto = true

pyproject.toml

@@ -23,7 +23,6 @@ classifiers = [
  "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dependencies = [
- "chardet>=4.0.0",
  "click>=8.0.0,<9",
  "pathspec>=0.9.0",
  "colorama>=0.4.6; sys_platform == 'win32'",

python/deptry/imports/extract.py

@@ -4,8 +4,7 @@
 from collections import defaultdict
 from typing import TYPE_CHECKING
 
-from deptry.imports.extractors import NotebookImportExtractor
-from deptry.rust import get_imports_from_py_files
+from deptry.rust import get_imports_from_ipynb_files, get_imports_from_py_files
 
 if TYPE_CHECKING:
  from pathlib import Path
@@ -19,7 +18,7 @@ def get_imported_modules_from_list_of_files(list_of_files: list[Path]) -> dict[s
  logging.info("Scanning %d %s...", len(list_of_files), "files" if len(list_of_files) > 1 else "file")
 
  py_files = [str(file) for file in list_of_files if file.suffix == ".py"]
- ipynb_files = [file for file in list_of_files if file.suffix == ".ipynb"]
+ ipynb_files = [str(file) for file in list_of_files if file.suffix == ".ipynb"]
 
  modules: dict[str, list[Location]] = defaultdict(list)
 
@@ -29,25 +28,17 @@ def get_imported_modules_from_list_of_files(list_of_files: list[Path]) -> dict[s
  for module, locations in convert_rust_locations_to_python_locations(rust_result).items():
  modules[module].extend(locations)
 
- # Process each .ipynb file individually
- for file in ipynb_files:
- for module, locations in get_imported_modules_from_ipynb_file(file).items():
+ # Process all .ipynb files in parallel using Rust
+ if ipynb_files:
+ rust_result = get_imports_from_ipynb_files(ipynb_files)
+ for module, locations in convert_rust_locations_to_python_locations(rust_result).items():
  modules[module].extend(locations)
 
  logging.debug("All imported modules: %s\n", modules)
 
  return modules
 
 
-def get_imported_modules_from_ipynb_file(path_to_file: Path) -> dict[str, list[Location]]:
- logging.debug("Scanning %s...", path_to_file)
-
- modules = NotebookImportExtractor(path_to_file).extract_imports()
-
- logging.debug("Found the following imports in %s: %s", path_to_file, modules)
- return modules
-
-
 def convert_rust_locations_to_python_locations(
  imported_modules: dict[str, list[RustLocation]],
 ) -> dict[str, list[Location]]:

python/deptry/rust.pyi

@@ -1,7 +1,7 @@
 from .rust import Location as RustLocation
 
 def get_imports_from_py_files(file_paths: list[str]) -> dict[str, list[RustLocation]]: ...
-def get_imports_from_py_file(file_path: str) -> dict[str, list[RustLocation]]: ...
+def get_imports_from_ipynb_files(file_paths: list[str]) -> dict[str, list[RustLocation]]: ...
 
 class Location:
  file: str

src/imports/ipynb.rs

@@ -0,0 +1,73 @@
+use crate::file_utils;
+use crate::location;
+
+use file_utils::read_file;
+use location::Location;
+use pyo3::exceptions::PySyntaxError;
+use pyo3::prelude::*;
+use pyo3::types::PyString;
+use rayon::prelude::*;
+use std::collections::HashMap;
+
+use super::shared;
+
+/// Processes multiple Python files in parallel to extract import statements and their locations.
+/// Accepts a list of file paths and returns a dictionary mapping module names to their import locations.
+#[pyfunction]
+pub fn get_imports_from_ipynb_files(py: Python, file_paths: Vec<&PyString>) -> PyResult<PyObject> {
+ let rust_file_paths: Vec<String> = file_paths
+ .iter()
+ .map(|py_str| py_str.to_str().unwrap().to_owned())
+ .collect();
+
+ let results: Vec<_> = rust_file_paths
+ .par_iter()
+ .map(|path_str| {
+ let result = _get_imports_from_ipynb_file(path_str);
+ shared::ThreadResult {
+ file: path_str.to_string(),
+ result,
+ }
+ })
+ .collect();
+
+ let (all_imports, errors) = shared::merge_results_from_threads(results);
+ shared::log_python_errors_as_warnings(&errors);
+
+ shared::convert_to_python_dict(py, all_imports)
+}
+
+/// Core helper function that extracts import statements and their locations from a single .ipynb file.
+/// Ensures robust error handling and provides clearer, more detailed comments.
+fn _get_imports_from_ipynb_file(path_str: &str) -> PyResult<HashMap<String, Vec<Location>>> {
+ let file_content = read_file(path_str)?;
+ let notebook: serde_json::Value =
+ serde_json::from_str(&file_content).map_err(|e| PySyntaxError::new_err(e.to_string()))?;
+ let cells = notebook["cells"]
+ .as_array()
+ .ok_or_else(|| PySyntaxError::new_err("Expected 'cells' to be an array"))?;
+ let python_code = _extract_code_from_notebook_cells(cells);
+
+ let ast = shared::get_ast_from_file_content(&python_code, path_str)?;
+ let imported_modules = shared::extract_imports_from_ast(ast);
+
+ Ok(shared::convert_imports_with_textranges_to_location_objects(
+ imported_modules,
+ path_str,
+ &python_code,
+ ))
+}
+
+/// Extracts and concatenates code from notebook code cells.
+fn _extract_code_from_notebook_cells(cells: &[serde_json::Value]) -> String {
+ let code_lines: Vec<String> = cells
+ .iter()
+ .filter(|cell| cell["cell_type"] == "code")
+ .flat_map(|cell| cell["source"].as_array())
+ .flatten()
+ .filter_map(|line| line.as_str())
+ .map(str::to_owned)
+ .collect();
+
+ code_lines.join("\n")
+}

src/imports/mod.rs

@@ -0,0 +1,3 @@
+pub mod ipynb;
+pub mod py;
+pub mod shared;