Skip to content

Commit

Permalink
Simplify import queries
Browse files Browse the repository at this point in the history
  • Loading branch information
@Peter554
Peter554 committed Dec 27, 2024
1 parent c48eb8e commit d4a7c4c
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 133 deletions.
147 changes: 43 additions & 104 deletions src/imports_info/queries/internal_imports.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use std::collections::{HashMap, HashSet};

use anyhow::Result;
use pathfinding::prelude::bfs_reach;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

use crate::{Error, ImportMetadata, ImportsInfo, PackageItemToken};

Expand Down Expand Up @@ -38,14 +37,12 @@ impl<'a> InternalImportsQueries<'a> {
) -> Result<HashSet<PackageItemToken>> {
self.imports_info.package_info.get_item(item)?;

self.for_every_package_item(item, |item| {
Ok(self
.imports_info
.internal_imports
.get(&item)
.unwrap()
.clone())
})
Ok(self
.imports_info
.internal_imports
.get(&item)
.unwrap()
.clone())
}

pub fn get_items_that_directly_import(
Expand All @@ -54,14 +51,12 @@ impl<'a> InternalImportsQueries<'a> {
) -> Result<HashSet<PackageItemToken>> {
self.imports_info.package_info.get_item(item)?;

self.for_every_package_item(item, |item| {
Ok(self
.imports_info
.reverse_internal_imports
.get(&item)
.unwrap()
.clone())
})
Ok(self
.imports_info
.reverse_internal_imports
.get(&item)
.unwrap()
.clone())
}

pub fn get_downstream_items(
Expand All @@ -70,16 +65,14 @@ impl<'a> InternalImportsQueries<'a> {
) -> Result<HashSet<PackageItemToken>> {
self.imports_info.package_info.get_item(item)?;

let mut items = self.for_every_package_item(item, |item| {
Ok(bfs_reach(item, |item| {
self.imports_info
.internal_imports
.get(item)
.unwrap()
.clone()
})
.collect())
})?;
let mut items = bfs_reach(item, |item| {
self.imports_info
.internal_imports
.get(item)
.unwrap()
.clone()
})
.collect::<HashSet<_>>();

items.remove(&item);

Expand All @@ -92,63 +85,20 @@ impl<'a> InternalImportsQueries<'a> {
) -> Result<HashSet<PackageItemToken>> {
self.imports_info.package_info.get_item(item)?;

let mut items = self.for_every_package_item(item, |item| {
Ok(bfs_reach(item, |item| {
self.imports_info
.reverse_internal_imports
.get(item)
.unwrap()
.clone()
})
.collect())
})?;
let mut items = bfs_reach(item, |item| {
self.imports_info
.reverse_internal_imports
.get(item)
.unwrap()
.clone()
})
.collect::<HashSet<_>>();

items.remove(&item);

Ok(items)
}

fn for_every_package_item<
F: Fn(PackageItemToken) -> Result<HashSet<PackageItemToken>> + Send + Sync,
>(
&'a self,
item: PackageItemToken,
get_items: F,
) -> Result<HashSet<PackageItemToken>> {
match item {
PackageItemToken::Package(package) => {
let package = self.imports_info.package_info.get_package(package)?;

let package_contents = self
.imports_info
.package_info
.get_package_contents(package.token);

let mut hs: HashSet<PackageItemToken> = HashSet::new();
hs.extend(
package_contents
.all_items
.par_iter()
.try_fold(
HashSet::new,
|mut hs, item| -> Result<HashSet<PackageItemToken>> {
hs.extend(get_items(*item)?);
Ok(hs)
},
)
.try_reduce(HashSet::new, |mut hs, v| {
hs.extend(v);
Ok(hs)
})?,
);
hs = &hs - &package_contents.all_items;

Ok(hs)
}
PackageItemToken::Module(_) => get_items(item),
}
}

pub fn get_import_metadata(
&'a self,
from: PackageItemToken,
Expand All @@ -160,6 +110,21 @@ impl<'a> InternalImportsQueries<'a> {
Err(Error::NoSuchImport)?
}
}

pub fn get_shortest_path(
&'a self,
from: PackageItemToken,
to: PackageItemToken,
) -> Result<Option<Vec<PackageItemToken>>> {
self.imports_info.package_info.get_item(from)?;
self.imports_info.package_info.get_item(to)?;

Ok(Some(vec![]))
}

pub fn path_exists(&'a self, from: PackageItemToken, to: PackageItemToken) -> Result<bool> {
Ok(self.get_shortest_path(from, to)?.is_some())
}
}

#[cfg(test)]
Expand Down Expand Up @@ -223,33 +188,16 @@ from . import red",
let package_info = PackageInfo::build(testpackage.path())?;
let imports_info = ImportsInfo::build(package_info)?;

let root_package = imports_info._item("testpackage");
let root_package_init = imports_info._item("testpackage.__init__");
let fruit = imports_info._item("testpackage.fruit");
let colors_package = imports_info._item("testpackage.colors");
let red = imports_info._item("testpackage.colors.red");

// A module
let imports = imports_info
.internal_imports()
.get_items_directly_imported_by(root_package_init)
.unwrap();
assert_eq!(imports, hashset! {fruit, red},);

// A package (removes internal items)
let imports = imports_info
.internal_imports()
.get_items_directly_imported_by(root_package)
.unwrap();
assert_eq!(imports, hashset! {},);

// Another package (removes internal items)
let imports = imports_info
.internal_imports()
.get_items_directly_imported_by(colors_package)
.unwrap();
assert_eq!(imports, hashset! {fruit},);

Ok(())
}

Expand Down Expand Up @@ -279,23 +227,14 @@ from testpackage import colors

let root_package_init = imports_info._item("testpackage.__init__");
let fruit = imports_info._item("testpackage.fruit");
let colors_package = imports_info._item("testpackage.colors");
let colors_package_init = imports_info._item("testpackage.colors.__init__");

// A module
let imports = imports_info
.internal_imports()
.get_items_that_directly_import(fruit)
.unwrap();
assert_eq!(imports, hashset! {root_package_init, colors_package_init},);

// A package (removes internal items)
let imports = imports_info
.internal_imports()
.get_items_that_directly_import(colors_package)
.unwrap();
assert_eq!(imports, hashset! {root_package_init, fruit},);

Ok(())
}

Expand Down
30 changes: 1 addition & 29 deletions src/package_info/queries.rs
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,11 @@
use anyhow::Result;
use std::{collections::HashSet, path::Path};

use maplit::hashset;
use std::path::Path;

use crate::package_info::{
Module, ModuleToken, Package, PackageInfo, PackageItem, PackageItemToken, PackageToken,
};
use crate::Error;

#[derive(Debug)]
pub(crate) struct PackageContents {
pub all_items: HashSet<PackageItemToken>,
}

impl PackageInfo {
pub fn get_item_by_path(&self, path: &Path) -> Option<PackageItem> {
if let Some(package) = self.packages_by_path.get(path) {
Expand Down Expand Up @@ -117,27 +110,6 @@ impl PackageInfo {
_ => None,
}
}

pub(crate) fn get_package_contents(&self, package: PackageToken) -> PackageContents {
let descendant_packages = self
.get_descendant_items(package)
.unwrap()
.filter_map(PackageInfo::filter_packages)
.map(|o| o.token.into())
.collect::<HashSet<_>>();

let descendant_modules = self
.get_descendant_items(package)
.unwrap()
.filter_map(PackageInfo::filter_modules)
.map(|o| o.token.into())
.collect::<HashSet<_>>();

let descendant_items = &descendant_packages | &descendant_modules;
let all_items = &hashset! {package.into()} | &descendant_items;

PackageContents { all_items }
}
}

#[cfg(test)]
Expand Down

0 comments on commit d4a7c4c

Please sign in to comment.