py_import_helper/utils/
formatting.rs

1//! Import formatting utilities
2//!
3//! This module provides functions for formatting Python import statements
4//! according to PEP 8 and common formatting standards (isort, Black).
5
6use super::parsing::custom_import_sort;
7use crate::types::{FormattingConfig, ImportStatement};
8use std::collections::{HashMap, HashSet};
9
10/// Format a list of imports, merging same-package imports where appropriate
11#[must_use]
12pub fn format_imports(imports: &[ImportStatement], config: &FormattingConfig) -> Vec<String> {
13    let mut package_imports: HashMap<String, Vec<&ImportStatement>> = HashMap::new();
14
15    // Group imports by package
16    for import in imports {
17        package_imports
18            .entry(import.package.clone())
19            .or_default()
20            .push(import);
21    }
22
23    let mut result = Vec::new();
24    let mut packages: Vec<_> = package_imports.keys().collect();
25    packages.sort();
26
27    for package in packages {
28        let imports_for_package = package_imports
29            .get(package)
30            .expect("BUG: package key must exist in HashMap");
31
32        if let Some(first) = imports_for_package.first() {
33            if imports_for_package.len() == 1 && first.items.is_empty() {
34                // Single direct import (e.g., "import os"), use as-is
35                result.push(first.statement.clone());
36            } else {
37                // Either multiple imports from same package, or a single import with items
38                // In both cases, apply formatting logic (may need multi-line)
39                result.extend(merge_package_imports(imports_for_package, config));
40            }
41        }
42    }
43
44    result
45}
46
47/// Merge multiple imports from the same package with configurable formatting
48#[must_use]
49pub fn merge_package_imports(
50    imports: &[&ImportStatement],
51    config: &FormattingConfig,
52) -> Vec<String> {
53    let mut all_items = HashSet::new();
54    let package = &imports[0].package;
55
56    // Collect all items being imported from this package
57    for import in imports {
58        all_items.extend(import.items.iter().cloned());
59    }
60
61    if all_items.is_empty() {
62        // Simple "import package" statements
63        return imports.iter().map(|i| i.statement.clone()).collect();
64    }
65
66    let mut sorted_items: Vec<_> = all_items.into_iter().collect();
67    sorted_items.sort_by(|a, b| custom_import_sort(a, b));
68
69    // Determine if we should use multi-line format
70    let should_use_multiline = if config.force_multiline {
71        true
72    } else if config.force_single_line {
73        false
74    } else {
75        // Auto-detect based on configuration
76        let total_chars = sorted_items.iter().map(String::len).sum::<usize>();
77        let import_line_length = "from ".len()
78            + package.len()
79            + " import ".len()
80            + total_chars
81            + (sorted_items.len() * 2);
82
83        sorted_items.len() >= config.multiline_threshold || import_line_length > config.line_length
84    };
85
86    if should_use_multiline {
87        // Multi-line with parentheses
88        let indent = " ".repeat(config.indent_size);
89        let mut result = vec![format!("from {} import (", package)];
90
91        for item in &sorted_items {
92            if config.use_trailing_comma {
93                result.push(format!("{}{},", indent, item));
94            } else {
95                result.push(format!("{}{}", indent, item));
96            }
97        }
98
99        result.push(")".to_string());
100        result
101    } else {
102        // Single line
103        vec![format!(
104            "from {} import {}",
105            package,
106            sorted_items.join(", ")
107        )]
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use crate::types::{ImportCategory, ImportType};
115
116    #[test]
117    fn test_merge_package_imports() {
118        let import1 = ImportStatement {
119            statement: "from typing import Any".to_string(),
120            category: ImportCategory::StandardLibrary,
121            import_type: ImportType::From,
122            package: "typing".to_string(),
123            items: vec!["Any".to_string()],
124            is_multiline: false,
125        };
126
127        let import2 = ImportStatement {
128            statement: "from typing import Optional".to_string(),
129            category: ImportCategory::StandardLibrary,
130            import_type: ImportType::From,
131            package: "typing".to_string(),
132            items: vec!["Optional".to_string()],
133            is_multiline: false,
134        };
135
136        let config = FormattingConfig::default();
137        let merged = merge_package_imports(&[&import1, &import2], &config);
138        assert_eq!(merged.len(), 1);
139        assert!(merged[0].contains("Any"));
140        assert!(merged[0].contains("Optional"));
141    }
142}
py_import_helper/utils/formatting.rs

py_import_helper/utils/
formatting.rs
