1use crate::fixtures::imports::is_stdlib_module;
23use crate::fixtures::string_utils::replace_identifier;
24use crate::fixtures::types::TypeImportSpec;
25use rustpython_parser::ast::{Mod, Stmt};
26use rustpython_parser::Mode;
27use std::collections::HashMap;
28use tracing::{debug, info, warn};
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ImportKind {
35 Future,
37 Stdlib,
39 ThirdParty,
41}
42
43#[derive(Debug, Clone, PartialEq)]
47pub struct ImportGroup {
48 pub first_line: usize,
50 pub last_line: usize,
52 pub kind: ImportKind,
54}
55
56#[derive(Debug, Clone, PartialEq)]
58pub struct ImportedName {
59 pub name: String,
61 pub alias: Option<String>,
63}
64
65impl ImportedName {
66 pub fn as_import_str(&self) -> String {
68 match &self.alias {
69 Some(alias) => format!("{} as {}", self.name, alias),
70 None => self.name.clone(),
71 }
72 }
73}
74
75#[derive(Debug, Clone, PartialEq)]
77pub struct ParsedFromImport {
78 pub line: usize,
80 pub end_line: usize,
82 pub module: String,
84 pub names: Vec<ImportedName>,
91 pub is_multiline: bool,
93}
94
95impl ParsedFromImport {
96 pub fn name_strings(&self) -> Vec<String> {
99 self.names.iter().map(|n| n.as_import_str()).collect()
100 }
101
102 pub fn has_star(&self) -> bool {
104 self.names.iter().any(|n| n.name == "*")
105 }
106}
107
108#[derive(Debug, Clone, PartialEq)]
110pub struct ParsedBareImport {
111 pub line: usize,
113 pub module: String,
115 pub alias: Option<String>,
117}
118
119#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121pub enum ParseSource {
122 Ast,
124 StringFallback,
126}
127
128pub struct ImportLayout {
132 pub groups: Vec<ImportGroup>,
134 pub from_imports: Vec<ParsedFromImport>,
136 #[allow(dead_code)]
142 pub bare_imports: Vec<ParsedBareImport>,
143 #[allow(dead_code)]
149 pub source: ParseSource,
150 lines: Vec<String>,
152}
153
154impl ImportLayout {
155 fn new(
156 groups: Vec<ImportGroup>,
157 from_imports: Vec<ParsedFromImport>,
158 bare_imports: Vec<ParsedBareImport>,
159 source: ParseSource,
160 content: &str,
161 ) -> Self {
162 let lines = content.lines().map(|l| l.to_string()).collect();
163 Self {
164 groups,
165 from_imports,
166 bare_imports,
167 source,
168 lines,
169 }
170 }
171
172 pub fn line_strs(&self) -> Vec<&str> {
174 self.lines.iter().map(|s| s.as_str()).collect()
175 }
176
177 pub fn line(&self, idx: usize) -> &str {
179 self.lines.get(idx).map(|s| s.as_str()).unwrap_or("")
180 }
181
182 pub fn find_matching_from_import(&self, module: &str) -> Option<&ParsedFromImport> {
189 self.from_imports
190 .iter()
191 .find(|fi| fi.module == module && !fi.has_star())
192 }
193}
194
195pub fn parse_import_layout(content: &str) -> ImportLayout {
208 match rustpython_parser::parse(content, Mode::Module, "") {
209 Ok(ast) => parse_layout_from_ast(&ast, content),
210 Err(e) => {
211 warn!("AST parse failed ({e}), using string fallback for import layout");
212 parse_layout_from_str(content)
213 }
214 }
215}
216
217pub fn classify_import_statement(statement: &str) -> ImportKind {
233 classify_module(top_level_module(statement).unwrap_or(""))
234}
235
236pub fn import_sort_key(name: &str) -> &str {
240 match name.find(" as ") {
241 Some(pos) => name[..pos].trim(),
242 None => name.trim(),
243 }
244}
245
246pub fn import_line_sort_key(line: &str) -> (u8, String) {
253 let trimmed = line.trim();
254 if let Some(rest) = trimmed.strip_prefix("import ") {
255 let module = rest.split_whitespace().next().unwrap_or("");
256 (0, module.to_lowercase())
257 } else if let Some(rest) = trimmed.strip_prefix("from ") {
258 let module = rest.split(" import ").next().unwrap_or("").trim();
259 (1, module.to_lowercase())
260 } else {
261 (2, String::new())
262 }
263}
264
265pub fn find_sorted_insert_position(
273 lines: &[&str],
274 group: &ImportGroup,
275 sort_key: &(u8, String),
276) -> u32 {
277 for (i, line) in lines
278 .iter()
279 .enumerate()
280 .take(group.last_line + 1)
281 .skip(group.first_line)
282 {
283 let existing_key = import_line_sort_key(line);
284 if *sort_key < existing_key {
285 return i as u32;
286 }
287 }
288 (group.last_line + 1) as u32
289}
290
291pub fn adapt_type_for_consumer(
313 return_type: &str,
314 fixture_imports: &[TypeImportSpec],
315 consumer_import_map: &HashMap<String, TypeImportSpec>,
316) -> (String, Vec<TypeImportSpec>) {
317 let mut adapted = return_type.to_string();
318 let mut remaining = Vec::new();
319
320 for spec in fixture_imports {
321 if spec.import_statement.starts_with("import ") {
322 let bare_module = spec
324 .import_statement
325 .strip_prefix("import ")
326 .unwrap()
327 .split(" as ")
328 .next()
329 .unwrap_or("")
330 .trim();
331
332 if bare_module.is_empty() {
333 remaining.push(spec.clone());
334 continue;
335 }
336
337 let prefix = format!("{}.", spec.check_name);
339 if !adapted.contains(&prefix) {
340 remaining.push(spec.clone());
341 continue;
342 }
343
344 let mut rewrites: Vec<(String, String)> = Vec::new(); let mut all_rewritable = true;
348 let mut pos = 0;
349
350 while let Some(hit) = adapted[pos..].find(&prefix) {
351 let abs = pos + hit;
352
353 if abs > 0 {
355 let prev = adapted.as_bytes()[abs - 1];
356 if prev.is_ascii_alphanumeric() || prev == b'_' {
357 pos = abs + prefix.len();
358 continue;
359 }
360 }
361
362 let name_start = abs + prefix.len();
363 let rest = &adapted[name_start..];
364 let name_end = rest
365 .find(|c: char| !c.is_alphanumeric() && c != '_')
366 .unwrap_or(rest.len());
367 let name = &rest[..name_end];
368
369 if name.is_empty() {
370 pos = name_start;
371 continue;
372 }
373
374 if let Some(consumer_spec) = consumer_import_map.get(name) {
375 let expected = format!("from {} import", bare_module);
376 if consumer_spec.import_statement.starts_with(&expected) {
377 let dotted = format!("{}.{}", spec.check_name, name);
378 if !rewrites.iter().any(|(d, _)| d == &dotted) {
379 rewrites.push((dotted, consumer_spec.check_name.clone()));
380 }
381 } else {
382 all_rewritable = false;
384 break;
385 }
386 } else {
387 all_rewritable = false;
389 break;
390 }
391
392 pos = name_start + name_end;
393 }
394
395 if all_rewritable && !rewrites.is_empty() {
396 for (dotted, short) in &rewrites {
397 adapted = adapted.replace(dotted.as_str(), short.as_str());
398 }
399 info!(
400 "Adapted type '{}' → '{}' (consumer already imports short names)",
401 return_type, adapted
402 );
403 } else {
404 debug!(
407 "adapt_type_for_consumer: cannot fully rewrite '{}' in '{}' \
408 (not all dotted names have matching from-imports in consumer) \
409 — keeping bare-import spec",
410 spec.check_name, return_type,
411 );
412 remaining.push(spec.clone());
413 }
414 } else if let Some((module, name_part)) = split_from_import(&spec.import_statement) {
415 let original_name = name_part.split(" as ").next().unwrap_or(name_part).trim();
424
425 if let Some(consumer_module_name) =
426 find_consumer_bare_import(consumer_import_map, module)
427 {
428 let dotted = format!("{}.{}", consumer_module_name, original_name);
429 let new_adapted = replace_identifier(&adapted, &spec.check_name, &dotted);
430 if new_adapted != adapted {
431 info!(
432 "Adapted type: '{}' → '{}' (consumer has bare import for '{}')",
433 spec.check_name, dotted, module
434 );
435 adapted = new_adapted;
436 } else {
438 remaining.push(spec.clone());
441 }
442 } else {
443 remaining.push(spec.clone());
444 }
445 } else {
446 remaining.push(spec.clone());
447 }
448 }
449
450 (adapted, remaining)
451}
452
453pub(crate) fn find_consumer_bare_import<'a>(
456 consumer_import_map: &'a HashMap<String, TypeImportSpec>,
457 module: &str,
458) -> Option<&'a str> {
459 for spec in consumer_import_map.values() {
460 if let Some(rest) = spec.import_statement.strip_prefix("import ") {
461 let module_part = rest.split(" as ").next().unwrap_or("").trim();
462 if module_part == module {
463 return Some(&spec.check_name);
464 }
465 }
466 }
467 None
468}
469
470pub(crate) fn can_merge_into(fi: &ParsedFromImport) -> bool {
479 !(fi.has_star() || fi.is_multiline && fi.names.is_empty())
480}
481
482fn classify_module(module: &str) -> ImportKind {
486 if module == "__future__" {
487 ImportKind::Future
488 } else if is_stdlib_module(module) {
489 ImportKind::Stdlib
490 } else {
491 ImportKind::ThirdParty
492 }
493}
494
495fn merge_kinds(a: ImportKind, b: ImportKind) -> ImportKind {
514 match (a, b) {
515 (ImportKind::Future, _) | (_, ImportKind::Future) => ImportKind::Future,
516 (ImportKind::ThirdParty, _) | (_, ImportKind::ThirdParty) => ImportKind::ThirdParty,
517 _ => ImportKind::Stdlib,
518 }
519}
520
521fn classify_import_line(line: &str) -> ImportKind {
532 let trimmed = line.trim();
533 if let Some(rest) = trimmed.strip_prefix("from ") {
534 let module = rest.split_whitespace().next().unwrap_or("");
536 classify_module(module.split('.').next().unwrap_or(""))
537 } else if let Some(rest) = trimmed.strip_prefix("import ") {
538 rest.split(',')
540 .filter_map(|part| {
541 let name = part.split_whitespace().next()?;
542 Some(classify_module(name.split('.').next().unwrap_or("")))
543 })
544 .fold(ImportKind::Stdlib, merge_kinds)
545 } else {
546 ImportKind::ThirdParty
547 }
548}
549
550fn top_level_module(line: &str) -> Option<&str> {
558 let trimmed = line.trim();
559 if let Some(rest) = trimmed.strip_prefix("from ") {
560 let module = rest.split_whitespace().next()?;
561 module.split('.').next()
562 } else if let Some(rest) = trimmed.strip_prefix("import ") {
563 let first = rest.split(',').next()?.trim();
564 let first = first.split_whitespace().next()?;
565 first.split('.').next()
566 } else {
567 None
568 }
569}
570
571fn split_from_import(statement: &str) -> Option<(&str, &str)> {
574 let rest = statement.strip_prefix("from ")?;
575 let (module, rest) = rest.split_once(" import ")?;
576 let module = module.trim();
577 let name = rest.trim();
578 if module.is_empty() || name.is_empty() {
579 None
580 } else {
581 Some((module, name))
582 }
583}
584
585fn parse_layout_from_ast(ast: &rustpython_parser::ast::Mod, content: &str) -> ImportLayout {
588 let line_starts = build_line_starts(content);
589 let offset_to_line = |offset: usize| -> usize {
590 line_starts
591 .partition_point(|&s| s <= offset)
592 .saturating_sub(1)
593 };
594
595 let mut from_imports: Vec<ParsedFromImport> = Vec::new();
596 let mut bare_imports: Vec<ParsedBareImport> = Vec::new();
597
598 let body = match ast {
599 Mod::Module(m) => &m.body,
600 _ => return parse_layout_from_str(content),
601 };
602
603 for stmt in body {
604 match stmt {
605 Stmt::ImportFrom(import_from) => {
606 let start_byte = import_from.range.start().to_usize();
607 let end_byte = import_from.range.end().to_usize();
608 let line = offset_to_line(start_byte);
609 let end_line = offset_to_line(end_byte.saturating_sub(1));
613
614 let mut module = import_from
615 .module
616 .as_ref()
617 .map(|m| m.to_string())
618 .unwrap_or_default();
619 if let Some(ref level) = import_from.level {
620 let level_val = level.to_usize();
621 if level_val > 0 {
622 let dots = ".".repeat(level_val);
623 module = dots + &module;
624 }
625 }
626
627 let names: Vec<ImportedName> = import_from
628 .names
629 .iter()
630 .map(|alias| ImportedName {
631 name: alias.name.to_string(),
632 alias: alias.asname.as_ref().map(|a| a.to_string()),
633 })
634 .collect();
635
636 let is_multiline = end_line > line;
637
638 from_imports.push(ParsedFromImport {
639 line,
640 end_line,
641 module,
642 names,
643 is_multiline,
644 });
645 }
646 Stmt::Import(import_stmt) => {
647 let start_byte = import_stmt.range.start().to_usize();
648 let line = offset_to_line(start_byte);
649 for alias in &import_stmt.names {
650 bare_imports.push(ParsedBareImport {
651 line,
652 module: alias.name.to_string(),
653 alias: alias.asname.as_ref().map(|a| a.to_string()),
654 });
655 }
656 }
657 _ => {}
658 }
659 }
660
661 let groups = build_groups_from_ast(&from_imports, &bare_imports);
662 ImportLayout::new(
663 groups,
664 from_imports,
665 bare_imports,
666 ParseSource::Ast,
667 content,
668 )
669}
670
671fn build_line_starts(content: &str) -> Vec<usize> {
675 let bytes = content.as_bytes();
676 let mut starts = vec![0usize];
677 for (i, &b) in bytes.iter().enumerate() {
678 if b == b'\n' {
679 starts.push(i + 1);
680 }
681 }
682 starts
683}
684
685struct ImportEvent {
688 first_line: usize,
689 last_line: usize,
690 top_module: String,
691}
692
693fn build_groups_from_ast(
701 from_imports: &[ParsedFromImport],
702 bare_imports: &[ParsedBareImport],
703) -> Vec<ImportGroup> {
704 let mut events: Vec<ImportEvent> = Vec::new();
705
706 for fi in from_imports {
707 let top = fi
708 .module
709 .trim_start_matches('.')
710 .split('.')
711 .next()
712 .unwrap_or("")
713 .to_string();
714 events.push(ImportEvent {
715 first_line: fi.line,
716 last_line: fi.end_line,
717 top_module: top,
718 });
719 }
720
721 for bi in bare_imports {
722 let top = bi.module.split('.').next().unwrap_or("").to_string();
728 events.push(ImportEvent {
729 first_line: bi.line,
730 last_line: bi.line,
731 top_module: top,
732 });
733 }
734
735 events.sort_by_key(|e| e.first_line);
736
737 let mut groups: Vec<ImportGroup> = Vec::new();
738 for event in events {
739 match groups.last_mut() {
740 Some(g) if event.first_line <= g.last_line + 1 => {
745 g.last_line = g.last_line.max(event.last_line);
746 g.kind = merge_kinds(g.kind, classify_module(&event.top_module));
747 }
748 _ => {
750 let kind = classify_module(&event.top_module);
751 groups.push(ImportGroup {
752 first_line: event.first_line,
753 last_line: event.last_line,
754 kind,
755 });
756 }
757 }
758 }
759
760 groups
761}
762
763fn parse_layout_from_str(content: &str) -> ImportLayout {
770 let lines: Vec<&str> = content.lines().collect();
771 let mut groups: Vec<ImportGroup> = Vec::new();
772 let mut from_imports: Vec<ParsedFromImport> = Vec::new();
773 let mut bare_imports: Vec<ParsedBareImport> = Vec::new();
774
775 let mut current_start: Option<usize> = None;
776 let mut current_last: usize = 0;
777 let mut current_kind = ImportKind::ThirdParty;
778 let mut seen_any_import = false;
779 let mut in_multiline = false;
780 let mut multiline_start: usize = 0;
781 let mut multiline_module: String = String::new();
782
783 for (i, &line) in lines.iter().enumerate() {
784 if in_multiline {
786 current_last = i;
787 let line_no_comment = line.split('#').next().unwrap_or("").trim_end();
788 if line_no_comment.contains(')') {
789 from_imports.push(ParsedFromImport {
790 line: multiline_start,
791 end_line: i,
792 module: multiline_module.clone(),
793 names: vec![],
796 is_multiline: true,
797 });
798 in_multiline = false;
799 }
800 continue;
801 }
802
803 if line.starts_with("import ") || line.starts_with("from ") {
805 seen_any_import = true;
806 if current_start.is_none() {
807 current_start = Some(i);
808 current_kind = classify_import_line(line);
809 } else {
810 current_kind = merge_kinds(current_kind, classify_import_line(line));
814 }
815 current_last = i;
816
817 if let Some(rest) = line.strip_prefix("from ") {
818 let module = rest
819 .split(" import ")
820 .next()
821 .unwrap_or("")
822 .trim()
823 .to_string();
824 let line_no_comment = line.split('#').next().unwrap_or("").trim_end();
825 if line_no_comment.contains('(') && !line_no_comment.contains(')') {
826 in_multiline = true;
828 multiline_start = i;
829 multiline_module = module;
830 } else {
831 if let Some(names_raw) = rest.split(" import ").nth(1) {
833 let names_str = names_raw.split('#').next().unwrap_or("").trim_end();
834 let names: Vec<ImportedName> = names_str
835 .split(',')
836 .filter_map(|n| {
837 let n = n.trim();
838 if n.is_empty() {
839 return None;
840 }
841 if let Some((name, alias)) = n.split_once(" as ") {
842 Some(ImportedName {
843 name: name.trim().to_string(),
844 alias: Some(alias.trim().to_string()),
845 })
846 } else {
847 Some(ImportedName {
848 name: n.to_string(),
849 alias: None,
850 })
851 }
852 })
853 .collect();
854 from_imports.push(ParsedFromImport {
855 line: i,
856 end_line: i,
857 module,
858 names,
859 is_multiline: false,
860 });
861 }
862 }
863 } else if let Some(rest) = line.strip_prefix("import ") {
864 for part in rest.split(',') {
866 let part = part.trim();
867 let (module_str, alias) = if let Some((m, a)) = part.split_once(" as ") {
868 (m.trim().to_string(), Some(a.trim().to_string()))
869 } else {
870 let m = part.split_whitespace().next().unwrap_or(part);
871 (m.to_string(), None)
872 };
873 if !module_str.is_empty() {
874 bare_imports.push(ParsedBareImport {
875 line: i,
876 module: module_str,
877 alias,
878 });
879 }
880 }
881 }
882 continue;
883 }
884
885 let trimmed = line.trim();
886
887 if trimmed.is_empty() || trimmed.starts_with('#') {
889 if let Some(start) = current_start.take() {
890 groups.push(ImportGroup {
891 first_line: start,
892 last_line: current_last,
893 kind: current_kind,
894 });
895 }
896 continue;
897 }
898
899 if seen_any_import {
901 if let Some(start) = current_start.take() {
902 groups.push(ImportGroup {
903 first_line: start,
904 last_line: current_last,
905 kind: current_kind,
906 });
907 }
908 break;
909 }
910 }
912
913 if let Some(start) = current_start {
915 groups.push(ImportGroup {
916 first_line: start,
917 last_line: current_last,
918 kind: current_kind,
919 });
920 }
921
922 ImportLayout::new(
923 groups,
924 from_imports,
925 bare_imports,
926 ParseSource::StringFallback,
927 content,
928 )
929}
930
931#[cfg(test)]
934mod tests {
935 use super::*;
936 use crate::fixtures::types::TypeImportSpec;
937 use std::collections::HashMap;
938
939 fn spec(check_name: &str, import_statement: &str) -> TypeImportSpec {
942 TypeImportSpec {
943 check_name: check_name.to_string(),
944 import_statement: import_statement.to_string(),
945 }
946 }
947
948 fn layout(lines: &[&str]) -> ImportLayout {
950 parse_import_layout(&lines.join("\n"))
951 }
952
953 #[test]
962 fn test_classify_future() {
963 assert_eq!(
964 classify_import_statement("from __future__ import annotations"),
965 ImportKind::Future
966 );
967 }
968
969 #[test]
970 fn test_classify_stdlib() {
971 assert_eq!(
972 classify_import_statement("from typing import Any"),
973 ImportKind::Stdlib
974 );
975 assert_eq!(
976 classify_import_statement("import pathlib"),
977 ImportKind::Stdlib
978 );
979 assert_eq!(
980 classify_import_statement("from collections.abc import Sequence"),
981 ImportKind::Stdlib
982 );
983 }
984
985 #[test]
986 fn test_classify_third_party() {
987 assert_eq!(
988 classify_import_statement("import pytest"),
989 ImportKind::ThirdParty
990 );
991 assert_eq!(
992 classify_import_statement("from myapp.db import Database"),
993 ImportKind::ThirdParty
994 );
995 }
996
997 #[test]
998 fn test_classify_comma_separated_stdlib() {
999 assert_eq!(
1001 classify_import_statement("import os, sys"),
1002 ImportKind::Stdlib
1003 );
1004 }
1005
1006 #[test]
1007 fn test_classify_comma_separated_mixed_kinds_first_module_wins() {
1008 assert_eq!(
1013 classify_import_statement("import os, pytest"),
1014 ImportKind::Stdlib );
1016 assert_eq!(
1017 classify_import_statement("import pytest, os"),
1018 ImportKind::ThirdParty );
1020 }
1021
1022 #[test]
1025 fn test_merge_kinds_future_wins_over_all() {
1026 assert_eq!(
1027 merge_kinds(ImportKind::Future, ImportKind::Stdlib),
1028 ImportKind::Future
1029 );
1030 assert_eq!(
1031 merge_kinds(ImportKind::Future, ImportKind::ThirdParty),
1032 ImportKind::Future
1033 );
1034 assert_eq!(
1035 merge_kinds(ImportKind::Stdlib, ImportKind::Future),
1036 ImportKind::Future
1037 );
1038 }
1039
1040 #[test]
1041 fn test_merge_kinds_third_party_wins_over_stdlib() {
1042 assert_eq!(
1043 merge_kinds(ImportKind::ThirdParty, ImportKind::Stdlib),
1044 ImportKind::ThirdParty
1045 );
1046 assert_eq!(
1047 merge_kinds(ImportKind::Stdlib, ImportKind::ThirdParty),
1048 ImportKind::ThirdParty
1049 );
1050 }
1051
1052 #[test]
1053 fn test_merge_kinds_same_kind_unchanged() {
1054 assert_eq!(
1055 merge_kinds(ImportKind::Stdlib, ImportKind::Stdlib),
1056 ImportKind::Stdlib
1057 );
1058 assert_eq!(
1059 merge_kinds(ImportKind::ThirdParty, ImportKind::ThirdParty),
1060 ImportKind::ThirdParty
1061 );
1062 }
1063
1064 #[test]
1067 fn test_classify_import_line_all_stdlib() {
1068 assert_eq!(classify_import_line("import os, sys"), ImportKind::Stdlib);
1069 }
1070
1071 #[test]
1072 fn test_classify_import_line_all_third_party() {
1073 assert_eq!(
1074 classify_import_line("import pytest, flask"),
1075 ImportKind::ThirdParty
1076 );
1077 }
1078
1079 #[test]
1080 fn test_classify_import_line_mixed_stdlib_first() {
1081 assert_eq!(
1083 classify_import_line("import os, pytest"),
1084 ImportKind::ThirdParty
1085 );
1086 }
1087
1088 #[test]
1089 fn test_classify_import_line_mixed_third_party_first() {
1090 assert_eq!(
1091 classify_import_line("import pytest, os"),
1092 ImportKind::ThirdParty
1093 );
1094 }
1095
1096 #[test]
1097 fn test_classify_import_line_three_modules_mixed() {
1098 assert_eq!(
1101 classify_import_line("import os, sys, pytest"),
1102 ImportKind::ThirdParty
1103 );
1104 }
1105
1106 #[test]
1107 fn test_classify_import_line_four_modules_stdlib_only() {
1108 assert_eq!(
1110 classify_import_line("import os, sys, re, pathlib"),
1111 ImportKind::Stdlib
1112 );
1113 }
1114
1115 #[test]
1116 fn test_classify_import_line_four_modules_third_party_last() {
1117 assert_eq!(
1119 classify_import_line("import os, sys, re, pytest"),
1120 ImportKind::ThirdParty
1121 );
1122 }
1123
1124 #[test]
1125 fn test_parse_groups_three_module_mixed_bare_import() {
1126 let l = layout(&["import os, sys, pytest", "", "def test(): pass"]);
1128 assert_eq!(l.groups.len(), 1);
1129 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1130 }
1131
1132 #[test]
1133 fn test_classify_import_line_from_import_unaffected() {
1134 assert_eq!(
1136 classify_import_line("from typing import Any"),
1137 ImportKind::Stdlib
1138 );
1139 assert_eq!(
1140 classify_import_line("from flask import Flask"),
1141 ImportKind::ThirdParty
1142 );
1143 }
1144
1145 #[test]
1148 fn test_parse_groups_mixed_bare_import_classified_as_third_party() {
1149 let l = layout(&["import os, pytest", "", "def test(): pass"]);
1152 assert_eq!(l.groups.len(), 1);
1153 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1154 }
1155
1156 #[test]
1157 fn test_parse_groups_mixed_bare_import_order_independent() {
1158 let l = layout(&["import pytest, os", "", "def test(): pass"]);
1160 assert_eq!(l.groups.len(), 1);
1161 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1162 }
1163
1164 #[test]
1165 fn test_parse_groups_all_stdlib_bare_import_unchanged() {
1166 let l = layout(&["import os, sys", "", "def test(): pass"]);
1167 assert_eq!(l.groups.len(), 1);
1168 assert_eq!(l.groups[0].kind, ImportKind::Stdlib);
1169 }
1170
1171 #[test]
1172 fn test_parse_groups_fallback_mixed_bare_import() {
1173 let l = parse_import_layout("import os, pytest\ndef test(:\n pass");
1175 assert_eq!(l.source, ParseSource::StringFallback);
1176 assert_eq!(l.groups.len(), 1);
1177 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1178 }
1179
1180 #[test]
1183 fn test_parse_layout_uses_ast_for_valid_python() {
1184 let l = layout(&["import os", "", "def test(): pass"]);
1185 assert_eq!(l.source, ParseSource::Ast);
1186 }
1187
1188 #[test]
1189 fn test_parse_layout_falls_back_for_invalid_python() {
1190 let l = parse_import_layout("import os\ndef test(:\n pass");
1191 assert_eq!(l.source, ParseSource::StringFallback);
1192 }
1193
1194 #[test]
1197 fn test_parse_groups_stdlib_and_third_party() {
1198 let l = layout(&[
1199 "import time",
1200 "",
1201 "import pytest",
1202 "from vcc.framework import fixture",
1203 "",
1204 "LOGGING_TIME = 2",
1205 ]);
1206 assert_eq!(l.groups.len(), 2);
1207 assert_eq!(l.groups[0].first_line, 0);
1208 assert_eq!(l.groups[0].last_line, 0);
1209 assert_eq!(l.groups[0].kind, ImportKind::Stdlib);
1210 assert_eq!(l.groups[1].first_line, 2);
1211 assert_eq!(l.groups[1].last_line, 3);
1212 assert_eq!(l.groups[1].kind, ImportKind::ThirdParty);
1213 }
1214
1215 #[test]
1216 fn test_parse_groups_single_third_party() {
1217 let l = layout(&["import pytest", "", "def test(): pass"]);
1218 assert_eq!(l.groups.len(), 1);
1219 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1220 assert_eq!(l.groups[0].first_line, 0);
1221 assert_eq!(l.groups[0].last_line, 0);
1222 }
1223
1224 #[test]
1225 fn test_parse_groups_no_imports() {
1226 let l = layout(&["def test(): pass"]);
1227 assert!(l.groups.is_empty());
1228 }
1229
1230 #[test]
1231 fn test_parse_groups_empty_file() {
1232 let l = layout(&[]);
1233 assert!(l.groups.is_empty());
1234 }
1235
1236 #[test]
1237 fn test_parse_groups_with_docstring_preamble() {
1238 let l = layout(&[
1239 r#""""Module docstring.""""#,
1240 "",
1241 "import pytest",
1242 "from pathlib import Path",
1243 "",
1244 "def test(): pass",
1245 ]);
1246 assert_eq!(l.groups.len(), 1);
1249 assert_eq!(l.groups[0].first_line, 2);
1250 assert_eq!(l.groups[0].last_line, 3);
1251 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1252 }
1253
1254 #[test]
1255 fn test_parse_groups_ignores_indented_imports() {
1256 let l = layout(&[
1257 "import pytest",
1258 "",
1259 "def test():",
1260 " from .utils import helper",
1261 " import os",
1262 ]);
1263 assert_eq!(l.groups.len(), 1);
1264 assert_eq!(l.groups[0].first_line, 0);
1265 assert_eq!(l.groups[0].last_line, 0);
1266 }
1267
1268 #[test]
1269 fn test_parse_groups_future_then_stdlib_then_third_party() {
1270 let l = layout(&[
1271 "from __future__ import annotations",
1272 "",
1273 "import os",
1274 "import time",
1275 "",
1276 "import pytest",
1277 "",
1278 "def test(): pass",
1279 ]);
1280 assert_eq!(l.groups.len(), 3);
1281 assert_eq!(l.groups[0].kind, ImportKind::Future);
1282 assert_eq!(l.groups[1].kind, ImportKind::Stdlib); assert_eq!(l.groups[2].kind, ImportKind::ThirdParty); }
1285
1286 #[test]
1287 fn test_parse_groups_with_comments_between() {
1288 let l = layout(&[
1289 "import os",
1290 "# stdlib above, third-party below",
1291 "import pytest",
1292 "",
1293 "def test(): pass",
1294 ]);
1295 assert_eq!(l.groups.len(), 2);
1297 assert_eq!(l.groups[0].kind, ImportKind::Stdlib);
1298 assert_eq!(l.groups[0].last_line, 0);
1299 assert_eq!(l.groups[1].kind, ImportKind::ThirdParty);
1300 assert_eq!(l.groups[1].first_line, 2);
1301 }
1302
1303 #[test]
1304 fn test_parse_groups_comma_separated_import_is_stdlib() {
1305 let l = layout(&[
1306 "import os, sys",
1307 "",
1308 "import pytest",
1309 "",
1310 "def test(): pass",
1311 ]);
1312 assert_eq!(l.groups.len(), 2);
1313 assert_eq!(l.groups[0].kind, ImportKind::Stdlib);
1314 assert_eq!(l.groups[0].first_line, 0);
1315 assert_eq!(l.groups[0].last_line, 0);
1316 assert_eq!(l.groups[1].kind, ImportKind::ThirdParty);
1317 }
1318
1319 #[test]
1320 fn test_parse_groups_multiline_import_single_group() {
1321 let l = layout(&["from liba import (", " moda,", " modb", ")"]);
1322 assert_eq!(l.groups.len(), 1);
1323 assert_eq!(l.groups[0].first_line, 0);
1324 assert_eq!(l.groups[0].last_line, 3);
1325 assert_eq!(l.groups[0].kind, ImportKind::ThirdParty);
1326 }
1327
1328 #[test]
1329 fn test_parse_groups_multiline_import_followed_by_third_party() {
1330 let l = layout(&[
1331 "from liba import (",
1332 " moda,",
1333 " modb",
1334 ")",
1335 "",
1336 "import pytest",
1337 "",
1338 "def test(): pass",
1339 ]);
1340 assert_eq!(l.groups.len(), 2);
1341 assert_eq!(l.groups[0].first_line, 0);
1342 assert_eq!(l.groups[0].last_line, 3);
1343 assert_eq!(l.groups[1].first_line, 5);
1344 assert_eq!(l.groups[1].last_line, 5);
1345 assert_eq!(l.groups[1].kind, ImportKind::ThirdParty);
1346 }
1347
1348 #[test]
1349 fn test_parse_groups_multiline_stdlib_then_third_party() {
1350 let l = layout(&[
1351 "from typing import (",
1352 " Any,",
1353 " Optional,",
1354 ")",
1355 "",
1356 "import pytest",
1357 "",
1358 "def test(): pass",
1359 ]);
1360 assert_eq!(l.groups.len(), 2);
1361 assert_eq!(l.groups[0].kind, ImportKind::Stdlib);
1362 assert_eq!(l.groups[0].first_line, 0);
1363 assert_eq!(l.groups[0].last_line, 3);
1364 assert_eq!(l.groups[1].kind, ImportKind::ThirdParty);
1365 assert_eq!(l.groups[1].first_line, 5);
1366 assert_eq!(l.groups[1].last_line, 5);
1367 }
1368
1369 #[test]
1370 fn test_parse_groups_inline_multiline_import() {
1371 let l = layout(&[
1372 "from typing import (Any,",
1373 " Optional)",
1374 "",
1375 "import pytest",
1376 ]);
1377 assert_eq!(l.groups.len(), 2);
1378 assert_eq!(l.groups[0].kind, ImportKind::Stdlib);
1379 assert_eq!(l.groups[0].first_line, 0);
1380 assert_eq!(l.groups[0].last_line, 1);
1381 assert_eq!(l.groups[1].kind, ImportKind::ThirdParty);
1382 assert_eq!(l.groups[1].first_line, 3);
1383 assert_eq!(l.groups[1].last_line, 3);
1384 }
1385
1386 #[test]
1389 fn test_from_imports_single_line() {
1390 let l = layout(&["from typing import Any, Optional"]);
1391 assert_eq!(l.from_imports.len(), 1);
1392 let fi = &l.from_imports[0];
1393 assert_eq!(fi.module, "typing");
1394 assert_eq!(fi.line, 0);
1395 assert_eq!(fi.end_line, 0);
1396 assert!(!fi.is_multiline);
1397 assert_eq!(fi.name_strings(), vec!["Any", "Optional"]);
1398 }
1399
1400 #[test]
1401 fn test_from_imports_with_alias() {
1402 let l = layout(&["from pathlib import Path as P"]);
1403 let fi = &l.from_imports[0];
1404 assert_eq!(fi.module, "pathlib");
1405 assert_eq!(fi.name_strings(), vec!["Path as P"]);
1406 }
1407
1408 #[test]
1409 fn test_from_imports_multiline_has_correct_end_line() {
1410 let l = layout(&["from typing import (", " Any,", " Optional,", ")"]);
1411 assert_eq!(l.from_imports.len(), 1);
1412 let fi = &l.from_imports[0];
1413 assert_eq!(fi.line, 0);
1414 assert_eq!(fi.end_line, 3);
1415 assert!(fi.is_multiline);
1416 if l.source == ParseSource::Ast {
1419 assert_eq!(fi.name_strings(), vec!["Any", "Optional"]);
1420 }
1421 }
1422
1423 #[test]
1424 fn test_bare_imports_comma_separated() {
1425 let l = layout(&["import os, sys"]);
1426 assert_eq!(l.bare_imports.len(), 2);
1427 assert_eq!(l.bare_imports[0].module, "os");
1428 assert_eq!(l.bare_imports[1].module, "sys");
1429 assert_eq!(l.bare_imports[0].line, 0);
1431 assert_eq!(l.bare_imports[1].line, 0);
1432 }
1433
1434 #[test]
1435 fn test_bare_import_with_alias() {
1436 let l = layout(&["import pathlib as pl"]);
1437 assert_eq!(l.bare_imports.len(), 1);
1438 assert_eq!(l.bare_imports[0].module, "pathlib");
1439 assert_eq!(l.bare_imports[0].alias, Some("pl".to_string()));
1440 }
1441
1442 #[test]
1445 fn test_find_matching_found() {
1446 let l = layout(&[
1447 "import pytest",
1448 "from typing import Optional",
1449 "",
1450 "def test(): pass",
1451 ]);
1452 let fi = l.find_matching_from_import("typing");
1453 assert!(fi.is_some());
1454 assert_eq!(fi.unwrap().name_strings(), vec!["Optional"]);
1455 }
1456
1457 #[test]
1458 fn test_find_matching_multiple_names() {
1459 let l = layout(&["from typing import Any, Optional, Union"]);
1460 let fi = l.find_matching_from_import("typing").unwrap();
1461 assert_eq!(fi.name_strings(), vec!["Any", "Optional", "Union"]);
1462 }
1463
1464 #[test]
1465 fn test_find_matching_not_found() {
1466 let l = layout(&["import pytest", "from pathlib import Path"]);
1467 assert!(l.find_matching_from_import("typing").is_none());
1468 }
1469
1470 #[test]
1471 fn test_find_matching_returns_multiline() {
1472 let l = layout(&["from typing import (", " Any,", " Optional,", ")"]);
1474 let fi = l.find_matching_from_import("typing");
1475 assert!(fi.is_some(), "multiline match should be returned");
1476 assert!(fi.unwrap().is_multiline);
1477 }
1478
1479 #[test]
1480 fn test_find_matching_skips_star() {
1481 let l = layout(&["from typing import *"]);
1482 assert!(l.find_matching_from_import("typing").is_none());
1483 }
1484
1485 #[test]
1486 fn test_find_matching_ignores_indented() {
1487 let l = layout(&[
1489 "import pytest",
1490 "",
1491 "def test():",
1492 " from typing import Any",
1493 ]);
1494 assert!(l.find_matching_from_import("typing").is_none());
1495 }
1496
1497 #[test]
1498 fn test_find_matching_with_inline_comment() {
1499 let l = layout(&["from typing import Any # comment"]);
1500 let fi = l.find_matching_from_import("typing").unwrap();
1501 assert_eq!(fi.name_strings(), vec!["Any"]);
1503 }
1504
1505 #[test]
1506 fn test_find_matching_aliases_preserved() {
1507 let l = layout(&["from os import path as p, getcwd as cwd"]);
1508 let fi = l.find_matching_from_import("os").unwrap();
1509 assert_eq!(fi.name_strings(), vec!["path as p", "getcwd as cwd"]);
1510 }
1511
1512 #[test]
1515 fn test_can_merge_single_line() {
1516 let fi = ParsedFromImport {
1517 line: 0,
1518 end_line: 0,
1519 module: "typing".to_string(),
1520 names: vec![ImportedName {
1521 name: "Any".to_string(),
1522 alias: None,
1523 }],
1524 is_multiline: false,
1525 };
1526 assert!(can_merge_into(&fi));
1527 }
1528
1529 #[test]
1530 fn test_can_merge_multiline_with_names() {
1531 let fi = ParsedFromImport {
1533 line: 0,
1534 end_line: 3,
1535 module: "typing".to_string(),
1536 names: vec![ImportedName {
1537 name: "Any".to_string(),
1538 alias: None,
1539 }],
1540 is_multiline: true,
1541 };
1542 assert!(can_merge_into(&fi));
1543 }
1544
1545 #[test]
1546 fn test_cannot_merge_multiline_without_names() {
1547 let fi = ParsedFromImport {
1549 line: 0,
1550 end_line: 3,
1551 module: "typing".to_string(),
1552 names: vec![],
1553 is_multiline: true,
1554 };
1555 assert!(!can_merge_into(&fi));
1556 }
1557
1558 #[test]
1559 fn test_cannot_merge_star() {
1560 let fi = ParsedFromImport {
1561 line: 0,
1562 end_line: 0,
1563 module: "typing".to_string(),
1564 names: vec![ImportedName {
1565 name: "*".to_string(),
1566 alias: None,
1567 }],
1568 is_multiline: false,
1569 };
1570 assert!(!can_merge_into(&fi));
1571 }
1572
1573 #[test]
1576 fn test_import_sort_key_plain() {
1577 assert_eq!(import_sort_key("Path"), "Path");
1578 }
1579
1580 #[test]
1581 fn test_import_sort_key_alias() {
1582 assert_eq!(import_sort_key("Path as P"), "Path");
1583 }
1584
1585 #[test]
1588 fn test_import_line_sort_key_bare_before_from() {
1589 let bare = import_line_sort_key("import os");
1590 let from = import_line_sort_key("from typing import Any");
1591 assert!(bare < from, "bare imports should sort before from-imports");
1592 }
1593
1594 #[test]
1595 fn test_import_line_sort_key_alphabetical_bare() {
1596 let a = import_line_sort_key("import os");
1597 let b = import_line_sort_key("import pathlib");
1598 let c = import_line_sort_key("import time");
1599 assert!(a < b);
1600 assert!(b < c);
1601 }
1602
1603 #[test]
1604 fn test_import_line_sort_key_alphabetical_from() {
1605 let a = import_line_sort_key("from pathlib import Path");
1606 let b = import_line_sort_key("from typing import Any");
1607 assert!(a < b);
1608 }
1609
1610 #[test]
1611 fn test_import_line_sort_key_dotted_module_ordering() {
1612 let short = import_line_sort_key("from vcc import conx_canoe");
1613 let long = import_line_sort_key("from vcc.conxtfw.framework import fixture");
1614 assert!(
1615 short < long,
1616 "shorter module path should sort before longer"
1617 );
1618 }
1619
1620 #[test]
1623 fn test_sorted_position_bare_before_existing_bare() {
1624 let lines = vec!["import os", "import time"];
1625 let group = ImportGroup {
1626 first_line: 0,
1627 last_line: 1,
1628 kind: ImportKind::Stdlib,
1629 };
1630 let key = import_line_sort_key("import pathlib");
1631 assert_eq!(find_sorted_insert_position(&lines, &group, &key), 1);
1632 }
1633
1634 #[test]
1635 fn test_sorted_position_from_after_all_bare() {
1636 let lines = vec!["import os", "import time"];
1637 let group = ImportGroup {
1638 first_line: 0,
1639 last_line: 1,
1640 kind: ImportKind::Stdlib,
1641 };
1642 let key = import_line_sort_key("from typing import Any");
1643 assert_eq!(find_sorted_insert_position(&lines, &group, &key), 2);
1644 }
1645
1646 #[test]
1647 fn test_sorted_position_from_between_existing_froms() {
1648 let lines = vec!["import pytest", "from aaa import X", "from zzz import Y"];
1649 let group = ImportGroup {
1650 first_line: 0,
1651 last_line: 2,
1652 kind: ImportKind::ThirdParty,
1653 };
1654 let key = import_line_sort_key("from mmm import Z");
1655 assert_eq!(find_sorted_insert_position(&lines, &group, &key), 2);
1656 }
1657
1658 #[test]
1659 fn test_sorted_position_before_everything() {
1660 let lines = vec!["import time", "from typing import Any"];
1661 let group = ImportGroup {
1662 first_line: 0,
1663 last_line: 1,
1664 kind: ImportKind::Stdlib,
1665 };
1666 let key = import_line_sort_key("import os");
1667 assert_eq!(find_sorted_insert_position(&lines, &group, &key), 0);
1668 }
1669
1670 #[test]
1673 fn test_adapt_dotted_to_short_when_consumer_has_from_import() {
1674 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1675 let mut consumer_map = HashMap::new();
1676 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1677 let (adapted, remaining) =
1678 adapt_type_for_consumer("pathlib.Path", &fixture_imports, &consumer_map);
1679 assert_eq!(adapted, "Path");
1680 assert!(
1681 remaining.is_empty(),
1682 "No import should remain: {:?}",
1683 remaining
1684 );
1685 }
1686
1687 #[test]
1688 fn test_adapt_no_rewrite_when_consumer_lacks_from_import() {
1689 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1690 let consumer_map = HashMap::new();
1691 let (adapted, remaining) =
1692 adapt_type_for_consumer("pathlib.Path", &fixture_imports, &consumer_map);
1693 assert_eq!(adapted, "pathlib.Path");
1694 assert_eq!(remaining.len(), 1);
1695 assert_eq!(remaining[0].import_statement, "import pathlib");
1696 }
1697
1698 #[test]
1699 fn test_adapt_no_rewrite_when_consumer_imports_from_different_module() {
1700 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1701 let mut consumer_map = HashMap::new();
1702 consumer_map.insert("Path".to_string(), spec("Path", "from mylib import Path"));
1703 let (adapted, remaining) =
1704 adapt_type_for_consumer("pathlib.Path", &fixture_imports, &consumer_map);
1705 assert_eq!(adapted, "pathlib.Path");
1706 assert_eq!(remaining.len(), 1);
1707 }
1708
1709 #[test]
1710 fn test_adapt_from_import_specs_pass_through_unchanged() {
1711 let fixture_imports = vec![spec("Path", "from pathlib import Path")];
1712 let consumer_map = HashMap::new();
1713 let (adapted, remaining) = adapt_type_for_consumer("Path", &fixture_imports, &consumer_map);
1714 assert_eq!(adapted, "Path");
1715 assert_eq!(remaining.len(), 1);
1716 assert_eq!(remaining[0].check_name, "Path");
1717 }
1718
1719 #[test]
1720 fn test_adapt_complex_generic_with_dotted_and_from() {
1721 let fixture_imports = vec![
1722 spec("Optional", "from typing import Optional"),
1723 spec("pathlib", "import pathlib"),
1724 ];
1725 let mut consumer_map = HashMap::new();
1726 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1727 consumer_map.insert(
1728 "Optional".to_string(),
1729 spec("Optional", "from typing import Optional"),
1730 );
1731 let (adapted, remaining) =
1732 adapt_type_for_consumer("Optional[pathlib.Path]", &fixture_imports, &consumer_map);
1733 assert_eq!(adapted, "Optional[Path]");
1734 assert_eq!(remaining.len(), 1);
1735 assert_eq!(remaining[0].check_name, "Optional");
1736 }
1737
1738 #[test]
1739 fn test_adapt_multiple_dotted_refs_same_module() {
1740 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1741 let mut consumer_map = HashMap::new();
1742 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1743 consumer_map.insert(
1744 "PurePath".to_string(),
1745 spec("PurePath", "from pathlib import PurePath"),
1746 );
1747 let (adapted, remaining) = adapt_type_for_consumer(
1748 "tuple[pathlib.Path, pathlib.PurePath]",
1749 &fixture_imports,
1750 &consumer_map,
1751 );
1752 assert_eq!(adapted, "tuple[Path, PurePath]");
1753 assert!(remaining.is_empty());
1754 }
1755
1756 #[test]
1757 fn test_adapt_partial_match_one_name_missing() {
1758 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1759 let mut consumer_map = HashMap::new();
1760 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1761 let (adapted, remaining) = adapt_type_for_consumer(
1762 "tuple[pathlib.Path, pathlib.PurePath]",
1763 &fixture_imports,
1764 &consumer_map,
1765 );
1766 assert_eq!(adapted, "tuple[pathlib.Path, pathlib.PurePath]");
1767 assert_eq!(remaining.len(), 1);
1768 }
1769
1770 #[test]
1771 fn test_adapt_aliased_bare_import() {
1772 let fixture_imports = vec![spec("pl", "import pathlib as pl")];
1773 let mut consumer_map = HashMap::new();
1774 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1775 let (adapted, remaining) =
1776 adapt_type_for_consumer("pl.Path", &fixture_imports, &consumer_map);
1777 assert_eq!(adapted, "Path");
1778 assert!(remaining.is_empty());
1779 }
1780
1781 #[test]
1782 fn test_adapt_no_false_match_on_prefix_substring() {
1783 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1784 let mut consumer_map = HashMap::new();
1785 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1786 let (adapted, remaining) =
1787 adapt_type_for_consumer("mypathlib.Path", &fixture_imports, &consumer_map);
1788 assert_eq!(adapted, "mypathlib.Path");
1789 assert_eq!(remaining.len(), 1);
1790 }
1791
1792 #[test]
1793 fn test_adapt_dotted_module_collections_abc() {
1794 let fixture_imports = vec![spec("collections.abc", "import collections.abc")];
1795 let mut consumer_map = HashMap::new();
1796 consumer_map.insert(
1797 "Iterable".to_string(),
1798 spec("Iterable", "from collections.abc import Iterable"),
1799 );
1800 let (adapted, remaining) = adapt_type_for_consumer(
1801 "collections.abc.Iterable[str]",
1802 &fixture_imports,
1803 &consumer_map,
1804 );
1805 assert_eq!(adapted, "Iterable[str]");
1806 assert!(remaining.is_empty());
1807 }
1808
1809 #[test]
1810 fn test_adapt_consumer_has_bare_import_no_rewrite() {
1811 let fixture_imports = vec![spec("pathlib", "import pathlib")];
1812 let mut consumer_map = HashMap::new();
1813 consumer_map.insert("pathlib".to_string(), spec("pathlib", "import pathlib"));
1814 let (adapted, remaining) =
1815 adapt_type_for_consumer("pathlib.Path", &fixture_imports, &consumer_map);
1816 assert_eq!(adapted, "pathlib.Path");
1817 assert_eq!(remaining.len(), 1);
1818 }
1819
1820 #[test]
1821 fn test_adapt_short_to_dotted_when_consumer_has_bare_import() {
1822 let fixture_imports = vec![spec("Path", "from pathlib import Path")];
1823 let mut consumer_map = HashMap::new();
1824 consumer_map.insert("pathlib".to_string(), spec("pathlib", "import pathlib"));
1825 let (adapted, remaining) = adapt_type_for_consumer("Path", &fixture_imports, &consumer_map);
1826 assert_eq!(adapted, "pathlib.Path");
1827 assert!(
1828 remaining.is_empty(),
1829 "No import should remain: {:?}",
1830 remaining
1831 );
1832 }
1833
1834 #[test]
1835 fn test_adapt_short_to_dotted_consumer_has_aliased_bare_import() {
1836 let fixture_imports = vec![spec("Path", "from pathlib import Path")];
1837 let mut consumer_map = HashMap::new();
1838 consumer_map.insert("pl".to_string(), spec("pl", "import pathlib as pl"));
1839 let (adapted, remaining) = adapt_type_for_consumer("Path", &fixture_imports, &consumer_map);
1840 assert_eq!(adapted, "pl.Path");
1841 assert!(remaining.is_empty());
1842 }
1843
1844 #[test]
1845 fn test_adapt_short_no_rewrite_when_consumer_lacks_bare_import() {
1846 let fixture_imports = vec![spec("Path", "from pathlib import Path")];
1847 let consumer_map = HashMap::new();
1848 let (adapted, remaining) = adapt_type_for_consumer("Path", &fixture_imports, &consumer_map);
1849 assert_eq!(adapted, "Path");
1850 assert_eq!(remaining.len(), 1);
1851 assert_eq!(remaining[0].check_name, "Path");
1852 }
1853
1854 #[test]
1855 fn test_adapt_short_to_dotted_generic_type() {
1856 let fixture_imports = vec![
1857 spec("Optional", "from typing import Optional"),
1858 spec("Path", "from pathlib import Path"),
1859 ];
1860 let mut consumer_map = HashMap::new();
1861 consumer_map.insert("pathlib".to_string(), spec("pathlib", "import pathlib"));
1862 let (adapted, remaining) =
1863 adapt_type_for_consumer("Optional[Path]", &fixture_imports, &consumer_map);
1864 assert_eq!(adapted, "Optional[pathlib.Path]");
1865 assert_eq!(remaining.len(), 1);
1866 assert_eq!(remaining[0].check_name, "Optional");
1867 }
1868
1869 #[test]
1870 fn test_adapt_short_to_dotted_word_boundary_safety() {
1871 let fixture_imports = vec![spec("Path", "from pathlib import Path")];
1872 let mut consumer_map = HashMap::new();
1873 consumer_map.insert("pathlib".to_string(), spec("pathlib", "import pathlib"));
1874 let (adapted, remaining) =
1875 adapt_type_for_consumer("PathLike", &fixture_imports, &consumer_map);
1876 assert_eq!(adapted, "PathLike");
1877 assert_eq!(remaining.len(), 1);
1878 assert_eq!(remaining[0].check_name, "Path");
1879 }
1880
1881 #[test]
1882 fn test_adapt_short_to_dotted_multiple_occurrences() {
1883 let fixture_imports = vec![spec("Path", "from pathlib import Path")];
1884 let mut consumer_map = HashMap::new();
1885 consumer_map.insert("pathlib".to_string(), spec("pathlib", "import pathlib"));
1886 let (adapted, remaining) =
1887 adapt_type_for_consumer("tuple[Path, Path]", &fixture_imports, &consumer_map);
1888 assert_eq!(adapted, "tuple[pathlib.Path, pathlib.Path]");
1889 assert!(remaining.is_empty());
1890 }
1891
1892 #[test]
1893 fn test_adapt_short_to_dotted_aliased_from_import() {
1894 let fixture_imports = vec![spec("P", "from pathlib import Path as P")];
1895 let mut consumer_map = HashMap::new();
1896 consumer_map.insert("pathlib".to_string(), spec("pathlib", "import pathlib"));
1897 let (adapted, remaining) = adapt_type_for_consumer("P", &fixture_imports, &consumer_map);
1898 assert_eq!(adapted, "pathlib.Path");
1899 assert!(remaining.is_empty());
1900 }
1901
1902 #[test]
1903 fn test_adapt_short_to_dotted_collections_abc() {
1904 let fixture_imports = vec![spec("Iterable", "from collections.abc import Iterable")];
1905 let mut consumer_map = HashMap::new();
1906 consumer_map.insert(
1907 "collections.abc".to_string(),
1908 spec("collections.abc", "import collections.abc"),
1909 );
1910 let (adapted, remaining) =
1911 adapt_type_for_consumer("Iterable[str]", &fixture_imports, &consumer_map);
1912 assert_eq!(adapted, "collections.abc.Iterable[str]");
1913 assert!(remaining.is_empty());
1914 }
1915
1916 #[test]
1917 fn test_adapt_both_directions_in_one_call() {
1918 let fixture_imports = vec![
1919 spec("Sequence", "from typing import Sequence"),
1920 spec("pathlib", "import pathlib"),
1921 ];
1922 let mut consumer_map = HashMap::new();
1923 consumer_map.insert("Path".to_string(), spec("Path", "from pathlib import Path"));
1924 consumer_map.insert("typing".to_string(), spec("typing", "import typing"));
1925 let (adapted, remaining) =
1926 adapt_type_for_consumer("Sequence[pathlib.Path]", &fixture_imports, &consumer_map);
1927 assert_eq!(adapted, "typing.Sequence[Path]");
1928 assert!(
1929 remaining.is_empty(),
1930 "Both specs should be dropped: {:?}",
1931 remaining
1932 );
1933 }
1934}