Skip to content

Commit

Permalink
Add implicit import metadata
Browse files Browse the repository at this point in the history
  • Loading branch information
@Peter554
Peter554 committed Dec 31, 2024
1 parent 300dc31 commit 08f98f4
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 63 deletions.
113 changes: 67 additions & 46 deletions src/imports_info/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,22 @@ use crate::{
Error, PackageItemIterator, Pypath,
};

/// An explicit import statement.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImportMetadata {
line_number: usize,
is_typechecking: bool,
pub struct ExplicitImportMetadata {
/// The line number of the import statement.
pub line_number: usize,
/// Whether the import statement is for typechecking only (`typing.TYPE_CHECKING`).
pub is_typechecking: bool,
}

/// Metadata associated with an import.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ImportMetadata {
/// An explicit import.
ExplicitImport(ExplicitImportMetadata),
/// An implicit import. E.g. all packages implicitly import their init modules.
ImplicitImport,
}

pub type PackageItemTokenSet = HashSet<PackageItemToken>;
Expand Down Expand Up @@ -95,7 +107,11 @@ impl ImportsInfo {
// By definition, packages import their init modules.
for package in package_info.get_all_items().filter_packages() {
if let Some(init_module) = package.init_module {
imports_info.add_internal_import(package.token.into(), init_module.into(), None)?;
imports_info.add_internal_import(
package.token.into(),
init_module.into(),
ImportMetadata::ImplicitImport,
)?;
}
}

Expand All @@ -105,10 +121,10 @@ impl ImportsInfo {
continue;
}

let metadata = ImportMetadata {
let metadata = ImportMetadata::ExplicitImport(ExplicitImportMetadata {
line_number: raw_import.line_number,
is_typechecking: raw_import.is_typechecking,
};
});

if package_info.pypath_is_internal(&raw_import.pypath)? {
let internal_item = {
Expand All @@ -131,9 +147,9 @@ impl ImportsInfo {
}
};

imports_info.add_internal_import(item, internal_item, Some(metadata))?;
imports_info.add_internal_import(item, internal_item, metadata)?;
} else if options.include_external_imports {
imports_info.add_external_import(item, raw_import.pypath, Some(metadata))?;
imports_info.add_external_import(item, raw_import.pypath, metadata)?;
}
}
}
Expand Down Expand Up @@ -178,28 +194,32 @@ impl ImportsInfo {
pub fn exclude_typechecking_imports(&self) -> Result<Self> {
let mut imports_info = self.clone();

let internal_imports =
self.internal_imports_metadata
.iter()
.filter_map(|((from, to), metadata)| {
let internal_imports = self.internal_imports_metadata.iter().filter_map(
|((from, to), metadata)| match metadata {
ImportMetadata::ExplicitImport(metadata) => {
if metadata.is_typechecking {
Some((*from, *to))
} else {
None
}
});
}
ImportMetadata::ImplicitImport => None,
},
);
imports_info = imports_info.exclude_internal_imports(internal_imports)?;

let external_imports =
self.external_imports_metadata
.iter()
.filter_map(|((from, to), metadata)| {
let external_imports = self.external_imports_metadata.iter().filter_map(
|((from, to), metadata)| match metadata {
ImportMetadata::ExplicitImport(metadata) => {
if metadata.is_typechecking {
Some((*from, to.clone()))
} else {
None
}
});
}
ImportMetadata::ImplicitImport => None,
},
);
imports_info = imports_info.exclude_external_imports(external_imports)?;

Ok(imports_info)
Expand All @@ -221,16 +241,14 @@ impl ImportsInfo {
&mut self,
from: PackageItemToken,
to: PackageItemToken,
metadata: Option<ImportMetadata>,
metadata: ImportMetadata,
) -> Result<()> {
self.internal_imports.entry(from).or_default().insert(to);
self.reverse_internal_imports
.entry(to)
.or_default()
.insert(from);
if let Some(metadata) = metadata {
self.internal_imports_metadata.insert((from, to), metadata);
}
self.internal_imports_metadata.insert((from, to), metadata);
Ok(())
}

Expand All @@ -256,15 +274,13 @@ impl ImportsInfo {
&mut self,
from: PackageItemToken,
to: Pypath,
metadata: Option<ImportMetadata>,
metadata: ImportMetadata,
) -> Result<()> {
self.external_imports
.entry(from)
.or_default()
.insert(to.clone());
if let Some(metadata) = metadata {
self.external_imports_metadata.insert((from, to), metadata);
}
self.external_imports_metadata.insert((from, to), metadata);
Ok(())
}

Expand Down Expand Up @@ -380,18 +396,19 @@ from django.db import models
assert_eq!(
imports_info.internal_imports_metadata,
hashmap! {
(root_package_init, a) => ImportMetadata{
(root_package, root_package_init) => ImportMetadata::ImplicitImport,
(root_package_init, a) => ImportMetadata::ExplicitImport(ExplicitImportMetadata {
line_number: 2,
is_typechecking: false,
},
(root_package_init, b) => ImportMetadata{
}),
(root_package_init, b) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 3,
is_typechecking: false,
},
(a, b) => ImportMetadata{
}),
(a, b) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 2,
is_typechecking: false,
}
})
}
);

Expand All @@ -408,10 +425,10 @@ from django.db import models
assert_eq!(
imports_info.external_imports_metadata,
hashmap! {
(b, "django.db.models".parse()?) => ImportMetadata{
(b, "django.db.models".parse()?) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 2,
is_typechecking: false,
},
}),
}
);

Expand Down Expand Up @@ -460,14 +477,15 @@ from testpackage import b
assert_eq!(
imports_info.internal_imports_metadata,
hashmap! {
(root_package_init, a) => ImportMetadata{
(root_package, root_package_init) => ImportMetadata::ImplicitImport,
(root_package_init, a) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 2,
is_typechecking: false,
},
(root_package_init, b) => ImportMetadata{
}),
(root_package_init, b) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 3,
is_typechecking: false,
},
}),
}
);

Expand Down Expand Up @@ -496,10 +514,11 @@ from testpackage import b
assert_eq!(
imports_info.internal_imports_metadata,
hashmap! {
(root_package_init, b) => ImportMetadata{
(root_package, root_package_init) => ImportMetadata::ImplicitImport,
(root_package_init, b) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 3,
is_typechecking: false,
},
}),
}
);

Expand Down Expand Up @@ -552,14 +571,15 @@ if TYPE_CHECKING:
assert_eq!(
imports_info.internal_imports_metadata,
hashmap! {
(root_package_init, a) => ImportMetadata{
(root_package, root_package_init) => ImportMetadata::ImplicitImport,
(root_package_init, a) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 4,
is_typechecking: false,
},
(root_package_init, b) => ImportMetadata{
}),
(root_package_init, b) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 7,
is_typechecking: true,
},
}),
}
);

Expand Down Expand Up @@ -588,10 +608,11 @@ if TYPE_CHECKING:
assert_eq!(
imports_info.internal_imports_metadata,
hashmap! {
(root_package_init, a) => ImportMetadata{
(root_package, root_package_init) => ImportMetadata::ImplicitImport,
(root_package_init, a) => ImportMetadata::ExplicitImport(ExplicitImportMetadata{
line_number: 4,
is_typechecking: false,
},
}),
}
);

Expand Down
9 changes: 5 additions & 4 deletions src/imports_info/queries/external_imports.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,14 @@ impl<'a> ExternalImportsQueries<'a> {
&'a self,
from: PackageItemToken,
to: T,
) -> Result<Option<&'a ImportMetadata>> {
) -> Result<&'a ImportMetadata> {
let to = to.into_pypath()?;
if self.direct_import_exists(from, to.borrow())? {
Ok(self
.imports_info
.external_imports_metadata
.get(&(from, to.borrow().clone())))
.get(&(from, to.borrow().clone()))
.unwrap())
} else {
Err(Error::NoSuchImport)?
}
Expand All @@ -71,7 +72,7 @@ mod tests {
use pretty_assertions::assert_eq;

use super::*;
use crate::{testpackage, testutils::TestPackage, PackageInfo};
use crate::{testpackage, testutils::TestPackage, ExplicitImportMetadata, PackageInfo};

#[test]
fn test_get_direct_imports() -> Result<()> {
Expand Down Expand Up @@ -137,7 +138,7 @@ mod tests {

assert_eq!(
metadata,
Some(&ImportMetadata {
&ImportMetadata::ExplicitImport(ExplicitImportMetadata {
line_number: 1,
is_typechecking: false
})
Expand Down
22 changes: 11 additions & 11 deletions src/imports_info/queries/internal_imports.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,13 @@ impl<'a> InternalImportsQueries<'a> {
&'a self,
from: PackageItemToken,
to: PackageItemToken,
) -> Result<Option<&'a ImportMetadata>> {
) -> Result<&'a ImportMetadata> {
if self.direct_import_exists(from, to)? {
Ok(self.imports_info.internal_imports_metadata.get(&(from, to)))
Ok(self
.imports_info
.internal_imports_metadata
.get(&(from, to))
.unwrap())
} else {
Err(Error::NoSuchImport)?
}
Expand Down Expand Up @@ -155,7 +159,7 @@ mod tests {
use pretty_assertions::assert_eq;

use super::*;
use crate::{testpackage, testutils::TestPackage, Error, PackageInfo};
use crate::{testpackage, testutils::TestPackage, Error, ExplicitImportMetadata, PackageInfo};

#[test]
fn test_get_direct_imports() -> Result<()> {
Expand Down Expand Up @@ -346,17 +350,13 @@ from testpackage import books",

let internal_imports = imports_info.internal_imports();

let metadata = internal_imports
.get_import_metadata(root_package, root_package_init)
.unwrap();
assert_eq!(metadata, None);
let metadata = internal_imports.get_import_metadata(root_package, root_package_init)?;
assert_eq!(metadata, &ImportMetadata::ImplicitImport);

let metadata = internal_imports
.get_import_metadata(root_package_init, fruit)
.unwrap();
let metadata = internal_imports.get_import_metadata(root_package_init, fruit)?;
assert_eq!(
metadata,
Some(&ImportMetadata {
&ImportMetadata::ExplicitImport(ExplicitImportMetadata {
line_number: 1,
is_typechecking: false
})
Expand Down
4 changes: 2 additions & 2 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ pub use testutils::TestPackage;

pub use errors::Error;
pub use imports_info::{
ExternalImportsQueries, ImportMetadata, ImportsInfo, ImportsInfoBuildOptions,
InternalImportsQueries, PackageItemTokenSet,
ExplicitImportMetadata, ExternalImportsQueries, ImportMetadata, ImportsInfo,
ImportsInfoBuildOptions, InternalImportsQueries, PackageItemTokenSet,
};
pub use package_info::{
Module, ModuleToken, Package, PackageInfo, PackageItem, PackageItemIterator, PackageItemToken,
Expand Down

0 comments on commit 08f98f4

Please sign in to comment.