py_import_helper/utils/
categorization.rs

1//! Import categorization utilities
2//!
3//! This module provides functions for categorizing Python imports into
4//! future, standard library, third-party, and local categories according to PEP 8.
5
6use super::parsing::extract_package;
7use crate::registry::constants::{COMMON_THIRD_PARTY_PACKAGES, PYTHON_STDLIB_MODULES};
8use crate::types::ImportCategory;
9use std::collections::HashSet;
10
11/// Categorize an import statement
12#[must_use]
13pub fn categorize_import<S: ::std::hash::BuildHasher>(
14    import_statement: &str,
15    local_package_prefixes: &HashSet<String, S>,
16) -> ImportCategory {
17    // Future imports always come first
18    if import_statement.starts_with("from __future__") {
19        return ImportCategory::Future;
20    }
21
22    let package = extract_package(import_statement);
23
24    // Determine category with priority order:
25    // 1. Local imports (relative or matching local prefixes)
26    // 2. Standard library (built-in or custom registered)
27    // 3. Third-party (custom registered or default)
28    if is_local_import(import_statement, local_package_prefixes) {
29        ImportCategory::Local
30    } else if is_standard_library_package(&package) {
31        ImportCategory::StandardLibrary
32    } else if is_common_third_party_package(&package) {
33        ImportCategory::ThirdParty
34    } else {
35        // Default to third-party for unknown packages
36        ImportCategory::ThirdParty
37    }
38}
39
40/// Check if this is a local/relative import
41#[must_use]
42pub fn is_local_import<S: ::std::hash::BuildHasher>(
43    import_statement: &str,
44    local_package_prefixes: &HashSet<String, S>,
45) -> bool {
46    // Check for relative imports
47    if import_statement.contains("from .")
48        || import_statement.contains("from ..")
49        || import_statement.contains("from ...")
50        || import_statement.contains("from ....")
51    {
52        return true;
53    }
54
55    let package = extract_package(import_statement);
56
57    // Check custom local package prefixes
58    for prefix in local_package_prefixes {
59        if package.starts_with(prefix.as_str()) {
60            return true;
61        }
62    }
63
64    false
65}
66
67/// Check if a package is part of Python's standard library
68#[must_use]
69pub fn is_standard_library_package(package: &str) -> bool {
70    PYTHON_STDLIB_MODULES.contains(&package)
71}
72
73/// Check if a package is a common third-party package
74#[must_use]
75pub fn is_common_third_party_package(package: &str) -> bool {
76    COMMON_THIRD_PARTY_PACKAGES.contains(&package)
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_categorize_future_import() {
85        let prefixes = HashSet::new();
86        let category = categorize_import("from __future__ import annotations", &prefixes);
87        assert_eq!(category, ImportCategory::Future);
88    }
89
90    #[test]
91    fn test_categorize_stdlib_import() {
92        let prefixes = HashSet::new();
93        let category = categorize_import("from typing import Any", &prefixes);
94        assert_eq!(category, ImportCategory::StandardLibrary);
95    }
96
97    #[test]
98    fn test_categorize_third_party_import() {
99        let prefixes = HashSet::new();
100        let category = categorize_import("from pydantic import BaseModel", &prefixes);
101        assert_eq!(category, ImportCategory::ThirdParty);
102    }
103
104    #[test]
105    fn test_categorize_local_import() {
106        let mut prefixes = HashSet::new();
107        prefixes.insert("myapp".to_string());
108
109        let category = categorize_import("from myapp.models import User", &prefixes);
110        assert_eq!(category, ImportCategory::Local);
111
112        let category = categorize_import("from .utils import helper", &prefixes);
113        assert_eq!(category, ImportCategory::Local);
114    }
115
116    #[test]
117    fn test_is_local_import() {
118        let mut prefixes = HashSet::new();
119        prefixes.insert("myapp".to_string());
120
121        assert!(is_local_import("from . import module", &prefixes));
122        assert!(is_local_import("from .. import parent", &prefixes));
123        assert!(is_local_import("from myapp.core import Engine", &prefixes));
124        assert!(!is_local_import("from typing import Any", &prefixes));
125    }
126
127    #[test]
128    fn test_is_standard_library_package() {
129        assert!(is_standard_library_package("typing"));
130        assert!(is_standard_library_package("os"));
131        assert!(is_standard_library_package("sys"));
132        assert!(is_standard_library_package("collections.abc"));
133        assert!(!is_standard_library_package("pydantic"));
134    }
135
136    #[test]
137    fn test_is_common_third_party_package() {
138        assert!(is_common_third_party_package("pydantic"));
139        assert!(is_common_third_party_package("httpx"));
140        assert!(is_common_third_party_package("pytest"));
141        assert!(!is_common_third_party_package("typing"));
142    }
143}
py_import_helper/utils/categorization.rs

py_import_helper/utils/
categorization.rs
