POC: Bulk rust import scanning

pyproject.toml

@@ -16,7 +16,6 @@ authors = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "joblib>=1.3.0",
     "typing-extensions>=3.10.0.0",
 ]
 classifiers = [

rust/src/lib.rs

@@ -8,6 +8,7 @@ use crate::errors::{GrimpError, GrimpResult};
 use crate::exceptions::{InvalidModuleExpression, ModuleNotPresent, NoSuchContainer, ParseError};
 use crate::graph::higher_order_queries::Level;
 use crate::graph::{Graph, Module, ModuleIterator, ModuleTokenIterator};
+use crate::import_parsing::ImportedObject;
 use crate::module_expressions::ModuleExpression;
 use derive_new::new;
 use itertools::Itertools;
@@ -17,11 +18,12 @@ use pyo3::prelude::*;
 use pyo3::types::{IntoPyDict, PyDict, PyFrozenSet, PyList, PySet, PyString, PyTuple};
 use rayon::prelude::*;
 use rustc_hash::FxHashSet;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
+use std::fs;
 
 #[pymodule]
 fn _rustgrimp(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_wrapped(wrap_pyfunction!(parse_imported_objects_from_code))?;
+    m.add_wrapped(wrap_pyfunction!(parse_imported_objects))?;
     m.add_class::<GraphWrapper>()?;
     m.add("ModuleNotPresent", py.get_type::<ModuleNotPresent>())?;
     m.add("NoSuchContainer", py.get_type::<NoSuchContainer>())?;
@@ -34,23 +36,52 @@ fn _rustgrimp(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
 }
 
 #[pyfunction]
-fn parse_imported_objects_from_code<'py>(
-    py: Python<'py>,
-    module_code: &str,
-) -> PyResult<Vec<Bound<'py, PyDict>>> {
-    let imports = import_parsing::parse_imports_from_code(module_code)?;
-
-    Ok(imports
+fn parse_imported_objects(
+    py: Python<'_>,
+    // [(module_name, module_path)]
+    modules: Vec<(String, String)>,
+) -> PyResult<HashMap<String, Vec<Bound<'_, PyDict>>>> {
+    let imported_objects = modules
+        .into_iter()
+        .par_bridge()
+        .try_fold(
+            HashMap::new,
+            |mut hm: HashMap<String, Vec<ImportedObject>>,
+             (module_name, module_path)|
+             -> GrimpResult<_> {
+                let module_code = fs::read_to_string(&module_path).unwrap();
+                let imports = import_parsing::parse_imports_from_code(&module_code)?;
+                hm.insert(module_name, imports);
+                Ok(hm)
+            },
+        )
+        .try_reduce(
+            HashMap::new,
+            |mut hm: HashMap<String, Vec<ImportedObject>>, chunk_hm| {
+                hm.extend(chunk_hm);
+                Ok(hm)
+            },
+        )?;
+
+    Ok(imported_objects
         .into_iter()
-        .map(|import| {
-            let dict = PyDict::new(py);
-            dict.set_item("name", import.name).unwrap();
-            dict.set_item("line_number", import.line_number).unwrap();
-            dict.set_item("line_contents", import.line_contents)
-                .unwrap();
-            dict.set_item("typechecking_only", import.typechecking_only)
-                .unwrap();
-            dict
+        .map(|(module_name, imported_objects)| {
+            (
+                module_name,
+                imported_objects
+                    .into_iter()
+                    .map(|import| {
+                        let dict = PyDict::new(py);
+                        dict.set_item("name", import.name).unwrap();
+                        dict.set_item("line_number", import.line_number).unwrap();
+                        dict.set_item("line_contents", import.line_contents)
+                            .unwrap();
+                        dict.set_item("typechecking_only", import.typechecking_only)
+                            .unwrap();
+                        dict
+                    })
+                    .collect(),
+            )
         })
         .collect())
 }

src/grimp/adaptors/importscanner.py

@@ -46,10 +46,9 @@ def scan_for_imports(
         """
         found_package = self._found_package_for_module(module)
         module_filename = self._determine_module_filename(module, found_package)
-        module_contents = self._read_module_contents(module_filename)
 
         try:
-            imported_objects = self._get_raw_imported_objects(module_contents)
+            imported_objects = self._get_raw_imported_objects(module.name, module_filename)
         except rust.ParseError as e:
             raise exceptions.SourceSyntaxError(
                 filename=module_filename,
@@ -137,8 +136,10 @@ def _module_is_package(self, module_filename: str) -> bool:
         return self.file_system.split(module_filename)[-1] == "__init__.py"
 
     @staticmethod
-    def _get_raw_imported_objects(module_contents: str) -> Set[_ImportedObject]:
-        imported_object_dicts = rust.parse_imported_objects_from_code(module_contents)
+    def _get_raw_imported_objects(module_name: str, module_filename: str) -> Set[_ImportedObject]:
+        imported_object_dicts = rust.parse_imported_objects([(module_name, module_filename)])[
+            module_name
+        ]
         return {_ImportedObject(**d) for d in imported_object_dicts}
 
     @staticmethod