Skip to content

Commit

Permalink
Document ImportsInfo
Browse files Browse the repository at this point in the history
  • Loading branch information
Peter554 committed Dec 31, 2024
1 parent 08f98f4 commit bd3e9b4
Show file tree
Hide file tree
Showing 8 changed files with 540 additions and 27 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ fn main() -> Result<()> {
This crate might be useful for something eventually, but right now it's mainly just
a hobby project for me to learn about rust.

If you are looking for something more mature, try [grimp](https://github.com/seddonym/grimp/) and/or [import-linter](https://github.com/seddonym/import-linter).
If you are looking for something more mature, try [grimp](https://github.com/seddonym/grimp/)/[import-linter](https://github.com/seddonym/import-linter).

## Limitations

Expand Down
2 changes: 1 addition & 1 deletion src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use std::path::PathBuf;
use rustpython_parser::ParseError;
use thiserror::Error;

use crate::{Pypath, ModuleToken, PackageToken};
use crate::{ModuleToken, PackageToken, Pypath};

#[allow(missing_docs)]
#[derive(Error, Debug, PartialEq)]
Expand Down
85 changes: 69 additions & 16 deletions src/imports_info/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ pub enum ImportMetadata {
ImplicitImport,
}

/// A set of [PackageItemToken]s.
pub type PackageItemTokenSet = HashSet<PackageItemToken>;

impl From<PackageItemToken> for PackageItemTokenSet {
Expand All @@ -39,6 +40,52 @@ impl From<PackageItemToken> for PackageItemTokenSet {
}
}

/// A rich representation of the imports within a python package.
///
/// ```
/// # use std::collections::HashSet;
/// # use anyhow::Result;
/// # use maplit::{hashmap, hashset};
/// # use pyimports::{testpackage,TestPackage,PackageInfo,ImportsInfo};
/// # fn main() -> Result<()> {
/// let test_package = testpackage! {
/// "__init__.py" => "from testpackage import a",
/// "a.py" => "from django.db import models"
/// };
///
/// let package_info = PackageInfo::build(test_package.path())?;
/// let imports_info = ImportsInfo::build(package_info)?;
///
/// let root_pkg = imports_info.package_info()
/// .get_item_by_pypath("testpackage")?.unwrap()
/// .token();
/// let root_init = imports_info.package_info()
/// .get_item_by_pypath("testpackage.__init__")?.unwrap()
/// .token();
/// let a = imports_info.package_info()
/// .get_item_by_pypath("testpackage.a")?.unwrap()
/// .token();
///
/// assert_eq!(
/// imports_info.internal_imports().get_direct_imports(),
/// hashmap! {
/// root_pkg => hashset!{root_init},
/// root_init => hashset!{a},
/// a => hashset!{},
/// }
/// );
///
/// assert_eq!(
/// imports_info.external_imports().get_direct_imports(),
/// hashmap! {
/// root_pkg => hashset!{},
/// root_init => hashset!{},
/// a => hashset!{"django.db.models".parse()?},
/// }
/// );
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone)]
pub struct ImportsInfo {
package_info: PackageInfo,
Expand All @@ -51,6 +98,7 @@ pub struct ImportsInfo {
external_imports_metadata: HashMap<(PackageItemToken, Pypath), ImportMetadata>,
}

/// Options for building an [ImportsInfo].
#[derive(Debug, Clone)]
pub struct ImportsInfoBuildOptions {
include_typechecking_imports: bool,
Expand All @@ -64,29 +112,34 @@ impl Default for ImportsInfoBuildOptions {
}

impl ImportsInfoBuildOptions {
/// Creates (default) build options.
pub fn new() -> Self {
ImportsInfoBuildOptions {
include_typechecking_imports: true,
include_external_imports: true,
}
}

/// Excludes typechecking imports (`typing.TYPE_CHECKING`).
pub fn exclude_typechecking_imports(mut self) -> Self {
self.include_typechecking_imports = false;
self
}

/// Excludes external imports.
pub fn exclude_external_imports(mut self) -> Self {
self.include_external_imports = false;
self
}
}

impl ImportsInfo {
/// Builds an [ImportsInfo] with the default options.
pub fn build(package_info: PackageInfo) -> Result<Self> {
ImportsInfo::build_with_options(package_info, ImportsInfoBuildOptions::new())
}

/// Builds an [ImportsInfo] with custom options.
pub fn build_with_options(
package_info: PackageInfo,
options: ImportsInfoBuildOptions,
Expand Down Expand Up @@ -135,7 +188,7 @@ impl ImportsInfo {
// An imported module.
item
} else if let Some(item) = package_info
.get_item_by_pypath(&raw_import.pypath.parent())?
.get_item_by_pypath(raw_import.pypath.parent())?
.map(|item| item.token())
{
// An imported module member.
Expand All @@ -157,40 +210,40 @@ impl ImportsInfo {
Ok(imports_info)
}

/// Returns a reference to the contained [PackageInfo].
pub fn package_info(&self) -> &PackageInfo {
&self.package_info
}

/// Returns an [InternalImportsQueries] object, that allows querying internal imports.
pub fn internal_imports(&self) -> InternalImportsQueries {
InternalImportsQueries { imports_info: self }
}

/// Returns an [ExternalImportsQueries] object, that allows querying external imports.
pub fn external_imports(&self) -> ExternalImportsQueries {
ExternalImportsQueries { imports_info: self }
}

pub fn exclude_internal_imports(
/// Excludes the passed imports.
/// Returns a new [ImportsInfo], leaves this instance unchanged.
pub fn exclude_imports(
&self,
imports: impl IntoIterator<Item = (PackageItemToken, PackageItemToken)>,
internal: impl IntoIterator<Item = (PackageItemToken, PackageItemToken)>,
external: impl IntoIterator<Item = (PackageItemToken, Pypath)>,
) -> Result<Self> {
let mut imports_info = self.clone();
for (from, to) in imports {
for (from, to) in internal {
imports_info.remove_internal_import(from, to)?;
}
Ok(imports_info)
}

pub fn exclude_external_imports(
&self,
imports: impl IntoIterator<Item = (PackageItemToken, Pypath)>,
) -> Result<Self> {
let mut imports_info = self.clone();
for (from, to) in imports {
for (from, to) in external {
imports_info.remove_external_import(from, to)?;
}
Ok(imports_info)
}

/// Excludes typechecking imports.
/// Returns a new [ImportsInfo], leaves this instance unchanged.
pub fn exclude_typechecking_imports(&self) -> Result<Self> {
let mut imports_info = self.clone();

Expand All @@ -206,7 +259,6 @@ impl ImportsInfo {
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)| match metadata {
Expand All @@ -220,7 +272,8 @@ impl ImportsInfo {
ImportMetadata::ImplicitImport => None,
},
);
imports_info = imports_info.exclude_external_imports(external_imports)?;

imports_info = imports_info.exclude_imports(internal_imports, external_imports)?;

Ok(imports_info)
}
Expand Down Expand Up @@ -489,7 +542,7 @@ from testpackage import b
}
);

let imports_info = imports_info.exclude_internal_imports(vec![(root_package_init, a)])?;
let imports_info = imports_info.exclude_imports(vec![(root_package_init, a)], vec![])?;

assert_eq!(
imports_info.internal_imports,
Expand Down
130 changes: 130 additions & 0 deletions src/imports_info/queries/external_imports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,85 @@ use anyhow::Result;

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

/// An object that allows querying external imports.
pub struct ExternalImportsQueries<'a> {
pub(crate) imports_info: &'a ImportsInfo,
}

impl<'a> ExternalImportsQueries<'a> {
/// Returns a map of all the direct imports.
///
/// ```
/// # use std::collections::HashSet;
/// # use anyhow::Result;
/// # use maplit::{hashmap, hashset};
/// # use pyimports::{testpackage,TestPackage,PackageInfo,ImportsInfo};
/// # fn main() -> Result<()> {
/// let test_package = testpackage! {
/// "__init__.py" => "from testpackage import a",
/// "a.py" => "from django.db import models"
/// };
///
/// let package_info = PackageInfo::build(test_package.path())?;
/// let imports_info = ImportsInfo::build(package_info)?;
///
/// let root_pkg = imports_info.package_info()
/// .get_item_by_pypath("testpackage")?.unwrap()
/// .token();
/// let root_init = imports_info.package_info()
/// .get_item_by_pypath("testpackage.__init__")?.unwrap()
/// .token();
/// let a = imports_info.package_info()
/// .get_item_by_pypath("testpackage.a")?.unwrap()
/// .token();
///
/// assert_eq!(
/// imports_info.external_imports().get_direct_imports(),
/// hashmap! {
/// root_pkg => hashset!{},
/// root_init => hashset!{},
/// a => hashset!{"django.db.models".parse()?},
/// }
/// );
/// # Ok(())
/// # }
/// ```
pub fn get_direct_imports(&self) -> HashMap<PackageItemToken, HashSet<Pypath>> {
self.imports_info.external_imports.clone()
}

/// Returns true if a direct import exists.
///
/// ```
/// # use std::collections::HashSet;
/// # use anyhow::Result;
/// # use maplit::{hashmap, hashset};
/// # use pyimports::{testpackage,TestPackage,PackageInfo,ImportsInfo};
/// # fn main() -> Result<()> {
/// let test_package = testpackage! {
/// "__init__.py" => "from testpackage import a",
/// "a.py" => "from django.db import models"
/// };
///
/// let package_info = PackageInfo::build(test_package.path())?;
/// let imports_info = ImportsInfo::build(package_info)?;
///
/// let root_init = imports_info.package_info()
/// .get_item_by_pypath("testpackage.__init__")?.unwrap()
/// .token();
/// let a = imports_info.package_info()
/// .get_item_by_pypath("testpackage.a")?.unwrap()
/// .token();
///
/// assert!(
/// imports_info.external_imports().direct_import_exists(a, "django.db.models")?,
/// );
/// assert!(
/// !imports_info.external_imports().direct_import_exists(root_init, "django.db.models")?,
/// );
/// # Ok(())
/// # }
/// ```
pub fn direct_import_exists<T: IntoPypath>(
&self,
from: PackageItemToken,
Expand All @@ -33,6 +103,36 @@ impl<'a> ExternalImportsQueries<'a> {
.contains(to.borrow()))
}

/// Returns the direct imports of the passed package item.
///
/// ```
/// # use std::collections::HashSet;
/// # use anyhow::Result;
/// # use maplit::{hashmap, hashset};
/// # use pyimports::{testpackage,TestPackage,PackageInfo,ImportsInfo};
/// # fn main() -> Result<()> {
/// let test_package = testpackage! {
/// "__init__.py" => "from testpackage import a",
/// "a.py" => "from django.db import models; import pydantic.BaseModel as BM"
/// };
///
/// let package_info = PackageInfo::build(test_package.path())?;
/// let imports_info = ImportsInfo::build(package_info)?;
///
/// let a = imports_info.package_info()
/// .get_item_by_pypath("testpackage.a")?.unwrap()
/// .token();
///
/// assert_eq!(
/// imports_info.external_imports().get_items_directly_imported_by(a)?,
/// hashset!{
/// "django.db.models".parse()?,
/// "pydantic.BaseModel".parse()?,
/// }
/// );
/// # Ok(())
/// # }
/// ```
pub fn get_items_directly_imported_by(
&'a self,
item: PackageItemToken,
Expand All @@ -47,6 +147,36 @@ impl<'a> ExternalImportsQueries<'a> {
.clone())
}

/// Returns the metadata associated with the passed import.
///
/// ```
/// # use std::collections::HashSet;
/// # use anyhow::Result;
/// # use maplit::{hashmap, hashset};
/// # use pyimports::{testpackage,TestPackage,PackageInfo,ImportsInfo,ImportMetadata,ExplicitImportMetadata};
/// # fn main() -> Result<()> {
/// let test_package = testpackage! {
/// "__init__.py" => "from testpackage import a",
/// "a.py" => "from django.db import models"
/// };
///
/// let package_info = PackageInfo::build(test_package.path())?;
/// let imports_info = ImportsInfo::build(package_info)?;
///
/// let a = imports_info.package_info()
/// .get_item_by_pypath("testpackage.a")?.unwrap()
/// .token();
///
/// assert_eq!(
/// imports_info.external_imports().get_import_metadata(a, "django.db.models")?,
/// &ImportMetadata::ExplicitImport(ExplicitImportMetadata {
/// line_number: 1,
/// is_typechecking: false
/// })
/// );
/// # Ok(())
/// # }
/// ```
pub fn get_import_metadata<T: IntoPypath>(
&'a self,
from: PackageItemToken,
Expand Down
Loading

0 comments on commit bd3e9b4

Please sign in to comment.