Add distill_external_module logic to handle namespace packages

Author: Peter554

rust/src/graph/builder/mod.rs

@@ -17,7 +17,7 @@ use cache::{ImportsCache, load_cache};
 
 mod utils;
 use utils::{
-    ResolvedImport, is_internal, is_package, path_to_module_name, resolve_external_module,
+    ResolvedImport, distill_external_module, is_internal, is_package, path_to_module_name,
     resolve_internal_module, resolve_relative_import,
 };
 
@@ -28,7 +28,7 @@ pub struct PackageSpec {
 }
 
 pub fn build_graph(
-    package: &PackageSpec,
+    package: &PackageSpec, // TODO(peter) Support multiple packages
     include_external_packages: bool,
     exclude_type_checking_imports: bool,
     cache_dir: Option<&PathBuf>,
@@ -48,6 +48,7 @@ pub fn build_graph(
         &parsed_modules,
         include_external_packages,
         exclude_type_checking_imports,
+        &HashSet::from([package.name.to_owned()]),
     );
 
     Ok(assemble_graph(&imports_by_module, &package.name))
@@ -269,6 +270,7 @@ fn resolve_imports(
     parsed_modules: &[ParsedModule],
     include_external_packages: bool,
     exclude_type_checking_imports: bool,
+    packages: &HashSet<String>,
 ) -> HashMap<String, HashSet<ResolvedImport>> {
     let all_modules: HashSet<String> = parsed_modules
         .iter()
@@ -304,14 +306,17 @@ fn resolve_imports(
                     line_contents: imported_object.line_contents.clone(),
                 });
             } else if include_external_packages {
-                // It's an external module and we're including them
-                let external_module = resolve_external_module(&absolute_import_name);
-                resolved_imports.insert(ResolvedImport {
-                    importer: parsed_module.module.name.to_string(),
-                    imported: external_module,
-                    line_number: imported_object.line_number,
-                    line_contents: imported_object.line_contents.clone(),
-                });
+                // Try to resolve as an external module
+                if let Some(external_module) =
+                    distill_external_module(&absolute_import_name, packages)
+                {
+                    resolved_imports.insert(ResolvedImport {
+                        importer: parsed_module.module.name.to_string(),
+                        imported: external_module,
+                        line_number: imported_object.line_number,
+                        line_contents: imported_object.line_contents.clone(),
+                    });
+                }
             }
         }
 

rust/src/graph/builder/utils.rs

@@ -20,12 +20,14 @@ pub fn is_package(module_path: &Path) -> bool {
         .unwrap_or(false)
 }
 
+/// Check if a module is a descendant of another module.
+pub fn is_descendant(module_name: &str, potential_ancestor: &str) -> bool {
+    module_name.starts_with(&format!("{}.", potential_ancestor))
+}
+
 /// Check if module is internal
 pub fn is_internal(module_name: &str, package: &str) -> bool {
-    if module_name == package || module_name.starts_with(&format!("{}.", package)) {
-        return true;
-    }
-    false
+    module_name == package || is_descendant(module_name, package)
 }
 
 /// Convert module path to module name
@@ -105,13 +107,62 @@ pub fn resolve_internal_module(
     None
 }
 
-/// Get external module name
-pub fn resolve_external_module(module_name: &str) -> String {
-    // For simplicity, just return the root module for external imports
-    // This matches the basic behavior from _distill_external_module
-    module_name
-        .split('.')
-        .next()
-        .unwrap_or(module_name)
-        .to_string()
+/// Given a module that we already know is external, turn it into a module to add to the graph.
+///
+/// The 'distillation' process involves removing any unwanted subpackages. For example,
+/// django.models.db should be turned into simply django.
+///
+/// The process is more complex for potential namespace packages, as it's not possible to
+/// determine the portion package simply from name. Rather than adding the overhead of a
+/// filesystem read, we just get the shallowest component that does not clash with an internal
+/// module namespace. Take, for example, foo.blue.alpha.one. If one of the found
+/// packages is foo.blue.beta, the module will be distilled to foo.blue.alpha.
+/// Alternatively, if the found package is foo.green, the distilled module will
+/// be foo.blue.
+///
+/// Returns None if the module is a parent of one of the internal packages (doesn't make sense,
+/// probably an import of a namespace package).
+pub fn distill_external_module(
+    module_name: &str,
+    found_package_names: &HashSet<String>,
+) -> Option<String> {
+    let module_root = module_name.split('.').next().unwrap();
+
+    for found_package in found_package_names {
+        // If it's a module that is a parent of the package, return None
+        // as it doesn't make sense and is probably an import of a namespace package.
+        if is_descendant(found_package, module_name) {
+            return None;
+        }
+    }
+
+    // If it shares a namespace with an internal module, get the shallowest component that does
+    // not clash with an internal module namespace.
+    let mut candidate_portions: Vec<String> = Vec::new();
+    let mut sorted_found_packages: Vec<&String> = found_package_names.iter().collect();
+    sorted_found_packages.sort();
+    sorted_found_packages.reverse();
+
+    for found_package in sorted_found_packages {
+        if is_descendant(found_package, module_root) {
+            let mut internal_components: Vec<&str> = found_package.split('.').collect();
+            let mut external_components: Vec<&str> = module_name.split('.').collect();
+            let mut external_namespace_components: Vec<&str> = vec![];
+            while external_components[0] == internal_components[0] {
+                external_namespace_components.push(external_components.remove(0));
+                internal_components.remove(0);
+            }
+            external_namespace_components.push(external_components[0]);
+            candidate_portions.push(external_namespace_components.join("."));
+        }
+    }
+
+    if !candidate_portions.is_empty() {
+        // If multiple internal modules share a namespace with this module, use the deepest one
+        // as we know that that will be a namespace too.
+        candidate_portions.sort_by_key(|portion| portion.split('.').count());
+        Some(candidate_portions.last().unwrap().clone())
+    } else {
+        Some(module_root.to_string())
+    }
 }