Skip to content

Commit

Permalink
Add squash_package
Browse files Browse the repository at this point in the history
  • Loading branch information
@Peter554
Peter554 committed Nov 26, 2023
1 parent 38d8b04 commit b318af4
Show file tree
Hide file tree
Showing 3 changed files with 203 additions and 3 deletions.
87 changes: 87 additions & 0 deletions src/import_graph/graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -415,4 +415,91 @@ impl ImportGraph {
reverse_imports,
})
}

pub fn squash_package(&self, package: &str) -> Result<ImportGraph> {
let package_to_squash = match self.packages_by_pypath.get(package) {
Some(package) => package,
None => {
return Err(Error::PackageNotFound(package.to_string()))?;
}
};
let binding = self
.child_modules(package)?
.into_iter()
.filter(|m| m.ends_with(".__init__"))
.collect::<Vec<_>>();
let init_module_pypath = binding.first().unwrap();
let init_module = self.modules_by_pypath.get(init_module_pypath).unwrap();

let packages_to_replace = self._descendant_packages(package)?;
let package_pypaths_to_replace = packages_to_replace
.iter()
.map(|p| p.pypath.clone())
.collect::<HashSet<_>>();
let modules_to_replace = {
let mut modules_to_replace = self._descendant_modules(package)?;
modules_to_replace.remove(init_module);
modules_to_replace
};
let module_pypaths_to_replace = modules_to_replace
.iter()
.map(|m| m.pypath.clone())
.collect::<HashSet<_>>();

let mut packages_by_pypath = self.packages_by_pypath.clone();
for pypath in packages_by_pypath.clone().keys() {
if package_pypaths_to_replace.contains(pypath) {
packages_by_pypath.remove(pypath);
}
}

let mut modules_by_pypath = self.modules_by_pypath.clone();
for pypath in modules_by_pypath.clone().keys() {
if module_pypaths_to_replace.contains(pypath) {
modules_by_pypath.remove(pypath);
}
}

let mut packages_by_module = self.packages_by_module.clone();
for (module, package) in packages_by_module.clone().iter() {
if modules_to_replace.contains(module) && packages_to_replace.contains(package) {
packages_by_module.remove(module);
packages_by_module.insert(Arc::clone(init_module), Arc::clone(package_to_squash));
} else if modules_to_replace.contains(module) {
packages_by_module.remove(module);
packages_by_module.insert(Arc::clone(init_module), Arc::clone(package));
} else if packages_to_replace.contains(package) {
packages_by_module.remove(module);
packages_by_module.insert(Arc::clone(module), Arc::clone(package_to_squash));
}
}

let mut imports = self.imports.clone();
for (module, imported_modules) in imports.clone().iter_mut() {
for imported_module in imported_modules.clone().iter() {
if modules_to_replace.contains(imported_module) {
imported_modules.remove(imported_module);
imported_modules.insert(Arc::clone(init_module));
}
}
imports.insert(Arc::clone(module), imported_modules.clone());
if modules_to_replace.contains(module) {
imports.remove(module);
imports
.get_mut(&Arc::clone(init_module))
.unwrap()
.extend(imported_modules.clone());
}
}

let reverse_imports = indexing::reverse_imports(&imports)?;

Ok(ImportGraph {
packages_by_pypath,
modules_by_pypath,
packages_by_module,
imports,
reverse_imports,
})
}
}
9 changes: 6 additions & 3 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,17 @@ fn main() -> Result<()> {
let root_package_path = Path::new(&args[1]);

let import_graph = {
let import_graph = ImportGraphBuilder::new(root_package_path).build()?;
if args.len() ==2 {
let mut import_graph = ImportGraphBuilder::new(root_package_path).build()?;
if args.len() == 2 {
import_graph
} else {
for child_package in import_graph.child_packages(&args[2])? {
import_graph = import_graph.squash_package(&child_package)?;
}
import_graph.subgraph(&args[2])?
}
};

let imports = import_graph.direct_imports_flat();
println!("source, target");
for (from_module, to_module) in imports {
Expand Down
110 changes: 110 additions & 0 deletions tests/test_import_graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1065,3 +1065,113 @@ fn test_subgraph() {
.collect()
);
}

#[test]
fn test_squash_package() {
let root_package_path = Path::new("./testpackages/somesillypackage");
let import_graph = ImportGraphBuilder::new(root_package_path).build().unwrap();
let squashed = import_graph
.squash_package("somesillypackage.child1")
.unwrap();
assert_eq!(
squashed.packages(),
hashset! {
"somesillypackage",
"somesillypackage.child1",
"somesillypackage.child2",
"somesillypackage.child3",
"somesillypackage.child4",
"somesillypackage.child5",
}
.into_iter()
.map(|s| s.to_string())
.collect()
);
assert_eq!(
squashed.modules(),
hashset! {
"somesillypackage.__init__",
"somesillypackage.a",
"somesillypackage.b",
"somesillypackage.c",
"somesillypackage.d",
"somesillypackage.e",
"somesillypackage.z",
"somesillypackage.child1.__init__",
"somesillypackage.child2.__init__",
"somesillypackage.child3.__init__",
"somesillypackage.child4.__init__",
"somesillypackage.child5.__init__",
}
.into_iter()
.map(|s| s.to_string())
.collect()
);
assert_eq!(
squashed.direct_imports(),
hashmap! {
"somesillypackage.__init__" => hashset!{
"somesillypackage.a",
"somesillypackage.b",
"somesillypackage.c",
"somesillypackage.d",
"somesillypackage.e",
"somesillypackage.child1.__init__",
"somesillypackage.child2.__init__",
"somesillypackage.child3.__init__",
"somesillypackage.child4.__init__",
"somesillypackage.child5.__init__",
},
"somesillypackage.a" => hashset!{
"somesillypackage.b",
"somesillypackage.c",
},
"somesillypackage.b" => hashset!{
"somesillypackage.c",
},
"somesillypackage.c" => hashset!{
"somesillypackage.d",
"somesillypackage.e",
},
"somesillypackage.d" => hashset!{
"somesillypackage.e"
},
"somesillypackage.e" => hashset!{},
"somesillypackage.z" => hashset! {
"somesillypackage.a",
"somesillypackage.b",
"somesillypackage.c",
"somesillypackage.d",
"somesillypackage.e",
"somesillypackage.child1.__init__",
"somesillypackage.child2.__init__",
"somesillypackage.child3.__init__",
"somesillypackage.child4.__init__",
"somesillypackage.child5.__init__",
},
"somesillypackage.child1.__init__" => hashset!{
"somesillypackage.a",
"somesillypackage.b",
"somesillypackage.c",
"somesillypackage.d",
"somesillypackage.e",
"somesillypackage.__init__",
"somesillypackage.child1.__init__",
"somesillypackage.child2.__init__",
"somesillypackage.child3.__init__",
"somesillypackage.child4.__init__",
"somesillypackage.child5.__init__",
},
"somesillypackage.child2.__init__" => hashset!{},
"somesillypackage.child3.__init__" => hashset!{},
"somesillypackage.child4.__init__" => hashset!{},
"somesillypackage.child5.__init__" => hashset!{},
}
.into_iter()
.map(|(k, v)| (
k.to_string(),
v.into_iter().map(|v| v.to_string()).collect()
))
.collect()
);
}

0 comments on commit b318af4

Please sign in to comment.