Skip to main content

pytest_language_server/fixtures/
import_analysis.rs

1//! Import layout analysis for Python files.
2//!
3//! This module is the single source of truth for import-related operations
4//! used across LSP providers (`code_action`, `inlay_hint`, and any future
5//! consumers).  It has no dependency on LSP types such as `TextEdit` — those
6//! belong in the provider layer.
7//!
8//! # Design
9//!
10//! [`parse_import_layout`] is the main entry point.  It tries to produce an
11//! [`ImportLayout`] via the `rustpython-parser` AST, which correctly handles
12//! multiline parenthesised imports, inline comments, and every other edge case
13//! that confounds a naive line-scanner.  If the file has syntax errors (common
14//! during active editing) it falls back to a simpler line-scan that mirrors
15//! the behaviour of the original string-based `parse_import_groups` function.
16//!
17//! [`adapt_type_for_consumer`] rewrites a fixture's return-type annotation
18//! string to match the consumer file's existing import style (dotted ↔ short).
19//! It is used by both `code_action` (which also inserts the necessary import
20//! statements) and `inlay_hint` (which only needs the display string).
21
22use 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// ─── Public types ─────────────────────────────────────────────────────────────
31
32/// Which isort / import-sort group an import belongs to.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ImportKind {
35    /// `from __future__ import …` — must appear before all other imports.
36    Future,
37    /// Python standard-library import.
38    Stdlib,
39    /// Third-party or project-local import.
40    ThirdParty,
41}
42
43/// A contiguous block of module-level import lines (separated from other
44/// blocks by blank lines or comment lines), with a classification for the
45/// whole group.
46#[derive(Debug, Clone, PartialEq)]
47pub struct ImportGroup {
48    /// 0-based index of the first line in this group.
49    pub first_line: usize,
50    /// 0-based index of the last line in this group (inclusive).
51    pub last_line: usize,
52    /// Classification based on the first import in the group.
53    pub kind: ImportKind,
54}
55
56/// One name (or alias) inside a `from X import …` statement.
57#[derive(Debug, Clone, PartialEq)]
58pub struct ImportedName {
59    /// The imported identifier, e.g. `"Path"` or `"*"`.
60    pub name: String,
61    /// The alias, if present, e.g. `"P"` for `Path as P`.
62    pub alias: Option<String>,
63}
64
65impl ImportedName {
66    /// Format as it appears in source: `"Name"` or `"Name as Alias"`.
67    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/// A parsed `from X import a, b` (or multiline variant) statement.
76#[derive(Debug, Clone, PartialEq)]
77pub struct ParsedFromImport {
78    /// 0-based line of the `from` keyword.
79    pub line: usize,
80    /// 0-based last line of the statement (equals `line` for single-line imports).
81    pub end_line: usize,
82    /// Fully-qualified module path, including leading dots for relative imports.
83    pub module: String,
84    /// Names imported.
85    ///
86    /// **Note**: for multiline imports parsed via the string fallback (i.e. the
87    /// file had syntax errors), this list will be empty because the individual
88    /// names cannot be reliably extracted line-by-line.  The AST path always
89    /// populates this field correctly.
90    pub names: Vec<ImportedName>,
91    /// `true` when the import spans multiple lines (parenthesised form).
92    pub is_multiline: bool,
93}
94
95impl ParsedFromImport {
96    /// Return each name formatted as an import string
97    /// (`"Name"` or `"Name as Alias"`).
98    pub fn name_strings(&self) -> Vec<String> {
99        self.names.iter().map(|n| n.as_import_str()).collect()
100    }
101
102    /// Whether this is a star-import (`from X import *`).
103    pub fn has_star(&self) -> bool {
104        self.names.iter().any(|n| n.name == "*")
105    }
106}
107
108/// A parsed `import X` or `import X as Y` statement.
109#[derive(Debug, Clone, PartialEq)]
110pub struct ParsedBareImport {
111    /// 0-based line of the `import` keyword.
112    pub line: usize,
113    /// The fully-qualified module name, e.g. `"pathlib"` or `"os.path"`.
114    pub module: String,
115    /// The alias, if present (`as Y`).
116    pub alias: Option<String>,
117}
118
119/// How the [`ImportLayout`] was derived — used for testing and diagnostics.
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121pub enum ParseSource {
122    /// Derived from a successfully parsed AST.
123    Ast,
124    /// Derived from simple line scanning (file has syntax errors).
125    StringFallback,
126}
127
128/// Complete import layout for the top-level import section of a Python file.
129///
130/// Obtained via [`parse_import_layout`].
131pub struct ImportLayout {
132    /// Classified groups of contiguous import lines.
133    pub groups: Vec<ImportGroup>,
134    /// All module-level `from X import …` statements, in source order.
135    pub from_imports: Vec<ParsedFromImport>,
136    /// All module-level `import X` statements, in source order.
137    ///
138    /// Exposed as part of the public API for future consumers; currently the
139    /// providers use the `existing_imports` `HashSet` for deduplication and
140    /// only read `from_imports` for merge decisions.
141    #[allow(dead_code)]
142    pub bare_imports: Vec<ParsedBareImport>,
143    /// How this layout was produced — `Ast` or `StringFallback`.
144    ///
145    /// Read in unit tests to assert which parser path was exercised; not
146    /// consumed by production code (the `warn!` in `parse_import_layout`
147    /// already records the fallback case in the log).
148    #[allow(dead_code)]
149    pub source: ParseSource,
150    /// Owned copy of each file line for `TextEdit` character-length lookups.
151    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    /// Borrow all lines as `&str` slices (for passing to sort/insert helpers).
173    pub fn line_strs(&self) -> Vec<&str> {
174        self.lines.iter().map(|s| s.as_str()).collect()
175    }
176
177    /// Get a single 0-based line as `&str` (empty string if out of bounds).
178    pub fn line(&self, idx: usize) -> &str {
179        self.lines.get(idx).map(|s| s.as_str()).unwrap_or("")
180    }
181
182    /// Find a module-level `from <module> import …` entry that is *not* a
183    /// star-import, regardless of whether it is single-line or multiline.
184    ///
185    /// Returns `None` if no matching entry exists or only star-imports match.
186    /// Used by `emit_kind_import_edits` to decide whether to merge a new name
187    /// into an existing import line or insert a fresh line.
188    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
195// ─── Public API ───────────────────────────────────────────────────────────────
196
197/// Parse the import layout of a Python source file.
198///
199/// Tries AST-based parsing first (via `rustpython-parser`); falls back to a
200/// simpler line-scan if the file has syntax errors (e.g. during active
201/// editing).
202///
203/// The returned [`ImportLayout`] contains:
204/// - `groups` — classified contiguous import blocks with line ranges,
205/// - `from_imports` / `bare_imports` — individual import statements,
206/// - `lines` — the file lines (for `TextEdit` character-length lookups).
207pub 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
217/// Classify an import statement string as [`ImportKind::Future`],
218/// [`ImportKind::Stdlib`], or [`ImportKind::ThirdParty`].
219///
220/// **Contract**: only the top-level package name of the *first* module in the
221/// statement is examined.  This is intentional: the function is called from
222/// `build_import_edits` exclusively with [`TypeImportSpec::import_statement`]
223/// values, which are always **single-module** strings such as
224/// `"from typing import Any"` or `"import pathlib"`.  Comma-separated
225/// multi-module lines (`import os, pytest`) can never reach this function
226/// through that path.
227///
228/// For group classification of raw file lines (the string-fallback parser) the
229/// same first-module heuristic is used, matching isort's own group-assignment
230/// behaviour.  In practice, mixed-kind lines are a style violation that tools
231/// like isort/ruff would split anyway.
232pub fn classify_import_statement(statement: &str) -> ImportKind {
233    classify_module(top_level_module(statement).unwrap_or(""))
234}
235
236/// Sort key for an imported name, stripping any `as alias` suffix.
237///
238/// `"Path as P"` → `"Path"`, `"Path"` → `"Path"`.
239pub 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
246/// Sort key for a full import line, following isort / ruff conventions:
247///
248/// 1. `import X` sorts **before** `from X import Y` (category `0` vs `1`).
249/// 2. Within each category, alphabetical by module path (case-insensitive).
250///
251/// Returns `(category, lowercased_module)`.
252pub 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
265/// Find the correct sorted insertion line within an existing import group,
266/// so that the result stays isort-sorted (bare before from, alphabetical by
267/// module within each sub-category).
268///
269/// Returns the 0-based line number for a point-insert.  When the new import
270/// sorts after every existing line in the group the position is
271/// `group.last_line + 1`.
272pub 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
291// ─── adapt_type_for_consumer ──────────────────────────────────────────────────
292
293/// Adapt a fixture's return-type annotation and import specs to the consumer
294/// file's existing import context.
295///
296/// Two adaptations are performed:
297///
298/// 1. **Dotted → short** — when a fixture uses `import pathlib` (bare) producing
299///    `pathlib.Path`, and the consumer already has `from pathlib import Path`,
300///    the annotation is shortened to `Path` and the bare-import spec is dropped.
301///
302/// 2. **Short → dotted** — when a fixture uses `from pathlib import Path`
303///    producing the short name `Path`, and the consumer already has
304///    `import pathlib` (bare), the annotation is lengthened to `pathlib.Path`
305///    and the from-import spec is dropped, respecting the consumer's style.
306///
307/// Returns `(adapted_type_string, remaining_import_specs)`.
308///
309/// Callers decide what to do with `remaining_import_specs`:
310/// - `code_action` inserts them as new import statements.
311/// - `inlay_hint` discards them (display only, no imports inserted).
312pub 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            // ── Case 1: bare-import spec → try dotted-to-short rewrite ───────
323            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            // Look for `check_name.Name` patterns in the type string.
338            let prefix = format!("{}.", spec.check_name);
339            if !adapted.contains(&prefix) {
340                remaining.push(spec.clone());
341                continue;
342            }
343
344            // Collect every `check_name.Name` occurrence and verify that the
345            // consumer already imports `Name` from the same module.
346            let mut rewrites: Vec<(String, String)> = Vec::new(); // (dotted, short)
347            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                // Guard against partial matches (e.g. `mypathlib.X` matching `pathlib.`).
354                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                        // Name imported from a different module — can't safely rewrite.
383                        all_rewritable = false;
384                        break;
385                    }
386                } else {
387                    // Name not in consumer's import map — can't rewrite.
388                    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                // Full-or-nothing: if any dotted name in the type string cannot
405                // be safely rewritten, keep the bare-import spec as-is.
406                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            // ── Case 2: from-import spec → try short-to-dotted rewrite ───────
416            //
417            // The fixture uses `from X import Y` so the type string contains
418            // the short name `Y`.  If the consumer already has `import X`
419            // (bare), we rewrite `Y` → `X.Y` and drop the from-import.
420
421            // Handle `from X import Y as Z` — the original name is `Y`, the
422            // check_name (used in the type string) is `Z`.
423            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                    // Drop the from-import spec — consumer's bare import covers it.
437                } else {
438                    // The check_name wasn't found as a standalone identifier in
439                    // the type string (word-boundary mismatch).  Keep the spec.
440                    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
453/// Find the `check_name` used by the consumer for a bare `import X` of the
454/// given module.  Returns `None` if the consumer does not have such an import.
455pub(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
470/// Returns `true` when we can safely merge new names into an existing
471/// [`ParsedFromImport`] — i.e. it is not a star-import, and it either has
472/// known names (AST path) or is single-line (both paths).
473///
474/// A multiline import from the string-fallback parser has `names.is_empty()`
475/// because individual names cannot be reliably extracted line-by-line from a
476/// file that failed to parse.  Merging into such an entry would lose existing
477/// names, so we fall back to inserting a new line instead.
478pub(crate) fn can_merge_into(fi: &ParsedFromImport) -> bool {
479    !(fi.has_star() || fi.is_multiline && fi.names.is_empty())
480}
481
482// ─── Internal helpers ─────────────────────────────────────────────────────────
483
484/// Classify a module name string as Future, Stdlib, or ThirdParty.
485fn 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
495/// Return the more restrictive of two import kinds.
496///
497/// Priority order: `Future` > `ThirdParty` > `Stdlib`.  This is a **binary
498/// reducer**, not the full algorithm.  The full N-module case is handled by
499/// composing this function with [`Iterator::fold`]:
500///
501/// ```text
502/// import os, sys, pytest, flask
503///   fold(Stdlib, merge_kinds):
504///     merge_kinds(Stdlib,      Stdlib)      // os      → Stdlib
505///     merge_kinds(Stdlib,      Stdlib)      // sys     → Stdlib
506///     merge_kinds(Stdlib,      ThirdParty)  // pytest  → ThirdParty
507///     merge_kinds(ThirdParty,  ThirdParty)  // flask   → ThirdParty
508///   result: ThirdParty  ✓
509/// ```
510///
511/// See [`classify_import_line`] for the call site that applies this to an
512/// arbitrary number of comma-separated modules.
513fn 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
521/// Classify an import **line** by inspecting *all* named modules, returning
522/// the most restrictive [`ImportKind`] found.
523///
524/// For a comma-separated bare import such as `import os, pytest`, every
525/// module is examined and `ThirdParty` wins over `Stdlib` (and `Future` wins
526/// over both).  This prevents a mixed line from being misclassified as
527/// `Stdlib` simply because the first module happens to be a stdlib package.
528///
529/// `from X import Y` lines have exactly one module, so the result is the
530/// same as calling [`classify_module`] directly.
531fn classify_import_line(line: &str) -> ImportKind {
532    let trimmed = line.trim();
533    if let Some(rest) = trimmed.strip_prefix("from ") {
534        // `from X import Y` — exactly one module.
535        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        // `import os, sys, pytest` — check every comma-separated module.
539        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
550/// Extract the top-level package name from the **first** module in an import
551/// line.  Only used for the `split_from_import` helper; all group-
552/// classification code should use [`classify_import_line`] instead.
553///
554/// - `"from collections.abc import X"` → `Some("collections")`
555/// - `"import os, sys"` → `Some("os")` (first module only)
556/// - `"x = 1"` → `None`
557fn 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
571/// Split `"from X import Y"` into `Some(("X", "Y"))`, or return `None` for
572/// bare `import X` statements and other non-matching strings.
573fn 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
585// ─── AST-based parser ─────────────────────────────────────────────────────────
586
587fn 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                // end() is exclusive; saturating_sub(1) gives the last byte of
610                // the statement, which is always on the same logical line as `)`
611                // for a multiline import.
612                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
671/// Build a byte-offset → 0-based-line-number lookup table from file content.
672///
673/// `result[i]` = byte offset of the start of line `i`.
674fn 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
685/// A lightweight event representing one import statement's line span and
686/// top-level module name, used only during group construction.
687struct ImportEvent {
688    first_line: usize,
689    last_line: usize,
690    top_module: String,
691}
692
693/// Group AST-derived import events into [`ImportGroup`]s.
694///
695/// Two imports are merged into the same group when the next import's first
696/// line is at most one greater than the previous import's last line (i.e. no
697/// blank line between them).  Comments in Python have no AST representation,
698/// so a comment between two imports creates a natural gap that is treated as a
699/// group boundary — matching the behaviour of the string-fallback parser.
700fn 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        // For `import os, sys` the AST yields two ParsedBareImport entries on
723        // the same line.  We emit an event for *every* entry so that the
724        // grouping step below can merge their kinds via `merge_kinds` — this
725        // correctly classifies `import os, pytest` as ThirdParty rather than
726        // Stdlib (first-module-wins).
727        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            // Adjacent or overlapping — extend the current group and update its
741            // kind to the most restrictive seen so far.  This handles both
742            // normal consecutive imports (different lines) and same-line
743            // comma-separated bare imports (same line, multiple events).
744            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            // Gap (blank line or comment) — start a new group.
749            _ => {
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
763// ─── String-fallback parser ───────────────────────────────────────────────────
764
765/// Line-scan fallback for files that fail to parse as valid Python.
766///
767/// Mirrors the behaviour of the original `parse_import_groups` function and
768/// additionally populates `from_imports` and `bare_imports`.
769fn 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        // ── Consume lines inside a parenthesised multiline import ────────────
785        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 cannot be reliably parsed from a fallback multiline;
794                    // see can_merge_into() for how callers handle this.
795                    names: vec![],
796                    is_multiline: true,
797                });
798                in_multiline = false;
799            }
800            continue;
801        }
802
803        // ── Module-level (unindented) import line ────────────────────────────
804        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                // A subsequent line in the same group: merge its kind so that a
811                // mixed line (e.g. `import os, pytest`) cannot downgrade the
812                // group's classification retroactively.
813                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                    // Opening of a multiline import.
827                    in_multiline = true;
828                    multiline_start = i;
829                    multiline_module = module;
830                } else {
831                    // Single-line from-import.
832                    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                // Bare imports (possibly comma-separated, e.g. `import os, sys`).
865                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        // ── Blank line or comment — close current group, keep scanning ───────
888        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        // ── Non-import, non-blank — stop if we've seen any import ────────────
900        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        // Before any import: preamble (module docstring, shebang, etc.) — keep going.
911    }
912
913    // Close final group if the file ends while still in an import section.
914    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// ─── Tests ────────────────────────────────────────────────────────────────────
932
933#[cfg(test)]
934mod tests {
935    use super::*;
936    use crate::fixtures::types::TypeImportSpec;
937    use std::collections::HashMap;
938
939    // ── helper ───────────────────────────────────────────────────────────
940
941    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    /// Build an ImportLayout from a slice of lines joined with newlines.
949    fn layout(lines: &[&str]) -> ImportLayout {
950        parse_import_layout(&lines.join("\n"))
951    }
952
953    // ── classify_import_statement ─────────────────────────────────────────
954    //
955    // classify_import_statement uses first-module-wins (intentional: it is
956    // only called from build_import_edits with single-module TypeImportSpec
957    // strings).  Group classification of raw file lines goes through
958    // classify_import_line, which checks *all* modules — see the
959    // parse_layout group tests below for that behaviour.
960
961    #[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        // Both modules are stdlib — unambiguous result.
1000        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        // classify_import_statement is called with single-module TypeImportSpec
1009        // strings only — it intentionally uses first-module-wins.  The mixed
1010        // case is handled correctly at the layout level by classify_import_line
1011        // (tested via test_parse_groups_mixed_bare_import_* below).
1012        assert_eq!(
1013            classify_import_statement("import os, pytest"),
1014            ImportKind::Stdlib // first-module-wins: os is stdlib
1015        );
1016        assert_eq!(
1017            classify_import_statement("import pytest, os"),
1018            ImportKind::ThirdParty // first-module-wins: pytest is third-party
1019        );
1020    }
1021
1022    // ── merge_kinds ───────────────────────────────────────────────────────
1023
1024    #[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    // ── classify_import_line ──────────────────────────────────────────────
1065
1066    #[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        // os is stdlib, pytest is third-party → ThirdParty wins regardless of order.
1082        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        // fold applies merge_kinds once per module — arbitrary length works.
1099        // os → Stdlib, sys → Stdlib, pytest → ThirdParty: result is ThirdParty.
1100        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        // All four stdlib → Stdlib.
1109        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        // Third-party appears last — fold must still catch it.
1118        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        // Three modules on one line; third-party is last → group is ThirdParty.
1127        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        // from-imports have exactly one module; behaviour is unchanged.
1135        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    // ── parse_import_layout — mixed bare imports ──────────────────────────
1146
1147    #[test]
1148    fn test_parse_groups_mixed_bare_import_classified_as_third_party() {
1149        // `import os, pytest` contains a third-party module — the group must
1150        // be ThirdParty regardless of which module appears first.
1151        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        // Result must be the same regardless of which module comes first.
1159        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        // Same assertion must hold for the string-fallback path.
1174        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    // ── parse_import_layout — parse source tracking ───────────────────────
1181
1182    #[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    // ── parse_import_layout — groups ──────────────────────────────────────
1195
1196    #[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        // pytest and pathlib are in the same contiguous block; group classified
1247        // by first import (pytest → ThirdParty).
1248        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); // os, time
1283        assert_eq!(l.groups[2].kind, ImportKind::ThirdParty); // pytest
1284    }
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        // Comment closes the first group, starts a new one.
1296        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    // ── parse_import_layout — from_imports / bare_imports fields ─────────
1387
1388    #[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        // AST path populates names; fallback path leaves them empty.
1417        // For valid Python (AST path), names must be present.
1418        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        // Both on line 0.
1430        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    // ── ImportLayout::find_matching_from_import ───────────────────────────
1443
1444    #[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        // New capability: multiline imports are now returned (not skipped).
1473        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        // Indented imports are not in module.body → not in from_imports.
1488        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        // Comment must NOT appear in name_strings.
1502        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    // ── can_merge_into ────────────────────────────────────────────────────
1513
1514    #[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        // AST path: multiline but names are populated → can merge.
1532        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        // String-fallback path: multiline with empty names → cannot merge.
1548        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    // ── import_sort_key ───────────────────────────────────────────────────
1574
1575    #[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    // ── import_line_sort_key ──────────────────────────────────────────────
1586
1587    #[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    // ── find_sorted_insert_position ───────────────────────────────────────
1621
1622    #[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    // ── adapt_type_for_consumer ───────────────────────────────────────────
1671
1672    #[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}