Generate ImportMap from module path to imported dependencies (#3243)

crates/ruff_cli/src/cache.rs

@@ -10,6 +10,7 @@ use path_absolutize::Absolutize;
 use ruff::message::Message;
 use ruff::settings::{flags, AllSettings, Settings};
 use ruff_cache::{CacheKey, CacheKeyHasher};
+use ruff_python_ast::imports::ImportMap;
 use serde::{Deserialize, Serialize};
 #[cfg(unix)]
 use std::os::unix::fs::PermissionsExt;
@@ -19,11 +20,13 @@ const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
 #[derive(Serialize)]
 struct CheckResultRef<'a> {
     messages: &'a [Message],
+    imports: &'a ImportMap,
 }
 
 #[derive(Deserialize)]
 struct CheckResult {
     messages: Vec<Message>,
+    imports: ImportMap,
 }
 
 fn content_dir() -> &'static Path {
@@ -95,14 +98,14 @@ pub fn get<P: AsRef<Path>>(
     metadata: &fs::Metadata,
     settings: &AllSettings,
     autofix: flags::Autofix,
-) -> Option<Vec<Message>> {
+) -> Option<(Vec<Message>, ImportMap)> {
     let encoded = read_sync(
         &settings.cli.cache_dir,
         cache_key(path, package, metadata, &settings.lib, autofix),
     )
     .ok()?;
     match bincode::deserialize::<CheckResult>(&encoded[..]) {
-        Ok(CheckResult { messages }) => Some(messages),
+        Ok(CheckResult { messages, imports }) => Some((messages, imports)),
         Err(e) => {
             error!("Failed to deserialize encoded cache entry: {e:?}");
             None
@@ -118,8 +121,9 @@ pub fn set<P: AsRef<Path>>(
     settings: &AllSettings,
     autofix: flags::Autofix,
     messages: &[Message],
+    imports: &ImportMap,
 ) {
-    let check_result = CheckResultRef { messages };
+    let check_result = CheckResultRef { messages, imports };
     if let Err(e) = write_sync(
         &settings.cli.cache_dir,
         cache_key(path, package, metadata, &settings.lib, autofix),

crates/ruff_cli/src/commands/run.rs

@@ -9,17 +9,18 @@ use log::{debug, error, warn};
 #[cfg(not(target_family = "wasm"))]
 use rayon::prelude::*;
 
-use crate::panic::catch_unwind;
 use ruff::message::{Location, Message};
 use ruff::registry::Rule;
 use ruff::resolver::PyprojectDiscovery;
 use ruff::settings::{flags, AllSettings};
 use ruff::{fs, packaging, resolver, warn_user_once, IOError, Range};
 use ruff_diagnostics::Diagnostic;
+use ruff_python_ast::imports::ImportMap;
 
 use crate::args::Overrides;
 use crate::cache;
 use crate::diagnostics::Diagnostics;
+use crate::panic::catch_unwind;
 
 /// Run the linter over a collection of files.
 pub fn run(
@@ -43,7 +44,7 @@ pub fn run(
 
     // Initialize the cache.
     if cache.into() {
-        fn init_cache(path: &std::path::Path) {
+        fn init_cache(path: &Path) {
             if let Err(e) = cache::init(path) {
                 error!("Failed to initialize cache at {}: {e:?}", path.display());
             }
@@ -115,15 +116,18 @@ pub fn run(
                     );
                     let settings = resolver.resolve(path, pyproject_strategy);
                     if settings.rules.enabled(Rule::IOError) {
-                        Diagnostics::new(vec![Message::from_diagnostic(
-                            Diagnostic::new(
-                                IOError { message },
-                                Range::new(Location::default(), Location::default()),
-                            ),
-                            format!("{}", path.display()),
-                            None,
-                            1,
-                        )])
+                        Diagnostics::new(
+                            vec![Message::from_diagnostic(
+                                Diagnostic::new(
+                                    IOError { message },
+                                    Range::new(Location::default(), Location::default()),
+                                ),
+                                format!("{}", path.display()),
+                                None,
+                                1,
+                            )],
+                            ImportMap::default(),
+                        )
                     } else {
                         Diagnostics::default()
                     }
@@ -137,7 +141,8 @@ pub fn run(
             acc += item;
             acc
         });
-
+    // TODO(chris): actually check the imports?
+    debug!("{:#?}", diagnostics.imports);
     diagnostics.messages.sort_unstable();
     let duration = start.elapsed();
     debug!("Checked {:?} files in: {:?}", paths.len(), duration);

crates/ruff_cli/src/diagnostics.rs

@@ -14,9 +14,10 @@ use similar::TextDiff;
 
 use ruff::fs;
 use ruff::jupyter::{is_jupyter_notebook, JupyterIndex, JupyterNotebook};
-use ruff::linter::{lint_fix, lint_only, FixTable, LinterResult};
+use ruff::linter::{lint_fix, lint_only, FixTable, FixerResult, LinterResult};
 use ruff::message::Message;
 use ruff::settings::{flags, AllSettings, Settings};
+use ruff_python_ast::imports::ImportMap;
 
 use crate::cache;
 
@@ -24,16 +25,18 @@ use crate::cache;
 pub struct Diagnostics {
     pub messages: Vec<Message>,
     pub fixed: FxHashMap<String, FixTable>,
+    pub imports: ImportMap,
     /// Jupyter notebook indexing table for each input file that is a jupyter notebook
     /// so we can rewrite the diagnostics in the end
     pub jupyter_index: FxHashMap<String, JupyterIndex>,
 }
 
 impl Diagnostics {
-    pub fn new(messages: Vec<Message>) -> Self {
+    pub fn new(messages: Vec<Message>, imports: ImportMap) -> Self {
         Self {
             messages,
             fixed: FxHashMap::default(),
+            imports,
             jupyter_index: FxHashMap::default(),
         }
     }
@@ -42,6 +45,7 @@ impl Diagnostics {
 impl AddAssign for Diagnostics {
     fn add_assign(&mut self, other: Self) {
         self.messages.extend(other.messages);
+        self.imports.extend(other.imports);
         for (filename, fixed) in other.fixed {
             if fixed.is_empty() {
                 continue;
@@ -58,7 +62,7 @@ impl AddAssign for Diagnostics {
 }
 
 /// Returns either an indexed python jupyter notebook or a diagnostic (which is empty if we skip)
-fn load_jupyter_notebook(path: &Path) -> Result<(String, JupyterIndex), Diagnostics> {
+fn load_jupyter_notebook(path: &Path) -> Result<(String, JupyterIndex), Box<Diagnostics>> {
     let notebook = match JupyterNotebook::read(path) {
         Ok(notebook) => {
             if !notebook
@@ -72,13 +76,13 @@ fn load_jupyter_notebook(path: &Path) -> Result<(String, JupyterIndex), Diagnost
                     "Skipping {} because it's not a Python notebook",
                     path.display()
                 );
-                return Err(Diagnostics::default());
+                return Err(Box::default());
             }
             notebook
         }
         Err(diagnostic) => {
             // Failed to read the jupyter notebook
-            return Err(Diagnostics {
+            return Err(Box::new(Diagnostics {
                 messages: vec![Message::from_diagnostic(
                     *diagnostic,
                     path.to_string_lossy().to_string(),
@@ -86,7 +90,7 @@ fn load_jupyter_notebook(path: &Path) -> Result<(String, JupyterIndex), Diagnost
                     1,
                 )],
                 ..Diagnostics::default()
-            });
+            }));
         }
     };
 
@@ -113,11 +117,11 @@ pub fn lint_path(
         && matches!(autofix, flags::FixMode::None | flags::FixMode::Generate)
     {
         let metadata = path.metadata()?;
-        if let Some(messages) =
+        if let Some((messages, imports)) =
             cache::get(path, package.as_ref(), &metadata, settings, autofix.into())
         {
             debug!("Cache hit for: {}", path.display());
-            return Ok(Diagnostics::new(messages));
+            return Ok(Diagnostics::new(messages, imports));
         }
         Some(metadata)
     } else {
@@ -130,7 +134,7 @@ pub fn lint_path(
     let (contents, jupyter_index) = if is_jupyter_notebook(path) {
         match load_jupyter_notebook(path) {
             Ok((contents, jupyter_index)) => (contents, Some(jupyter_index)),
-            Err(diagnostics) => return Ok(diagnostics),
+            Err(diagnostics) => return Ok(*diagnostics),
         }
     } else {
         (std::fs::read_to_string(path)?, None)
@@ -139,13 +143,16 @@ pub fn lint_path(
     // Lint the file.
     let (
         LinterResult {
-            data: messages,
+            data: (messages, imports),
             error: parse_error,
         },
         fixed,
     ) = if matches!(autofix, flags::FixMode::Apply | flags::FixMode::Diff) {
-        if let Ok((result, transformed, fixed)) =
-            lint_fix(&contents, path, package, noqa, &settings.lib)
+        if let Ok(FixerResult {
+            result,
+            transformed,
+            fixed,
+        }) = lint_fix(&contents, path, package, noqa, &settings.lib)
         {
             if !fixed.is_empty() {
                 if matches!(autofix, flags::FixMode::Apply) {
@@ -187,6 +194,8 @@ pub fn lint_path(
         (result, fixed)
     };
 
+    let imports = imports.unwrap_or_default();
+
     if let Some(err) = parse_error {
         // Notify the user of any parse errors.
         error!(
@@ -210,6 +219,7 @@ pub fn lint_path(
                 settings,
                 autofix.into(),
                 &messages,
+                &imports,
             );
         }
     }
@@ -231,6 +241,7 @@ pub fn lint_path(
     Ok(Diagnostics {
         messages,
         fixed: FxHashMap::from_iter([(fs::relativize_path(path), fixed)]),
+        imports,
         jupyter_index,
     })
 }
@@ -248,12 +259,16 @@ pub fn lint_stdin(
     // Lint the inputs.
     let (
         LinterResult {
-            data: messages,
+            data: (messages, imports),
             error: parse_error,
         },
         fixed,
     ) = if matches!(autofix, flags::FixMode::Apply | flags::FixMode::Diff) {
-        if let Ok((result, transformed, fixed)) = lint_fix(
+        if let Ok(FixerResult {
+            result,
+            transformed,
+            fixed,
+        }) = lint_fix(
             contents,
             path.unwrap_or_else(|| Path::new("-")),
             package,
@@ -312,6 +327,8 @@ pub fn lint_stdin(
         (result, fixed)
     };
 
+    let imports = imports.unwrap_or_default();
+
     if let Some(err) = parse_error {
         error!(
             "Failed to parse {}: {err}",
@@ -325,6 +342,7 @@ pub fn lint_stdin(
             fs::relativize_path(path.unwrap_or_else(|| Path::new("-"))),
             fixed,
         )]),
+        imports,
         jupyter_index: FxHashMap::default(),
     })
 }
@@ -341,7 +359,7 @@ mod test {
         // No diagnostics is used as skip signal
         assert_eq!(
             load_jupyter_notebook(path).unwrap_err(),
-            Diagnostics::default()
+            Box::<Diagnostics>::default()
         );
     }
 }