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

crates/ruff_python_ast/src/imports.rs

@@ -1,4 +1,8 @@
-use std::fmt;
+use rustc_hash::FxHashMap;
+use rustpython_parser::ast::Location;
+use serde::{Deserialize, Serialize};
+
+use crate::types::Range;
 
 /// A representation of an individual name imported via any import statement.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -38,8 +42,8 @@ impl<'a> Import<'a> {
     }
 }
 
-impl fmt::Display for AnyImport<'_> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl std::fmt::Display for AnyImport<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
             AnyImport::Import(import) => write!(f, "{import}"),
             AnyImport::ImportFrom(import_from) => write!(f, "{import_from}"),
@@ -47,8 +51,8 @@ impl fmt::Display for AnyImport<'_> {
     }
 }
 
-impl fmt::Display for Import<'_> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl std::fmt::Display for Import<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         write!(f, "import {}", self.name.name)?;
         if let Some(as_name) = self.name.as_name {
             write!(f, " as {as_name}")?;
@@ -57,8 +61,8 @@ impl fmt::Display for Import<'_> {
     }
 }
 
-impl fmt::Display for ImportFrom<'_> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl std::fmt::Display for ImportFrom<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         write!(f, "from ")?;
         if let Some(level) = self.level {
             write!(f, "{}", ".".repeat(level))?;
@@ -70,3 +74,59 @@ impl fmt::Display for ImportFrom<'_> {
         Ok(())
     }
 }
+
+/// A representation of a module reference in an import statement.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ModuleImport {
+    module: String,
+    location: Location,
+    end_location: Location,
+}
+
+impl ModuleImport {
+    pub fn new(module: String, location: Location, end_location: Location) -> Self {
+        Self {
+            module,
+            location,
+            end_location,
+        }
+    }
+}
+
+impl From<&ModuleImport> for Range {
+    fn from(import: &ModuleImport) -> Range {
+        Range::new(import.location, import.end_location)
+    }
+}
+
+/// A representation of the import dependencies between modules.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ImportMap {
+    /// A map from dot-delimited module name to the list of imports in that module.
+    module_to_imports: FxHashMap<String, Vec<ModuleImport>>,
+}
+
+impl ImportMap {
+    pub fn new() -> Self {
+        Self {
+            module_to_imports: FxHashMap::default(),
+        }
+    }
+
+    pub fn insert(&mut self, module: String, imports_vec: Vec<ModuleImport>) {
+        self.module_to_imports.insert(module, imports_vec);
+    }
+
+    pub fn extend(&mut self, other: Self) {
+        self.module_to_imports.extend(other.module_to_imports);
+    }
+}
+
+impl<'a> IntoIterator for &'a ImportMap {
+    type Item = (&'a String, &'a Vec<ModuleImport>);
+    type IntoIter = std::collections::hash_map::Iter<'a, String, Vec<ModuleImport>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.module_to_imports.iter()
+    }
+}