kimun 0.34.1

//! Python import extraction and resolution.
//!
//! An absolute import names a module under a root of the search path, which
//! the source does not tell; a relative one, a module beside the importer.
//! Both are kept, and the resolver leaves out what the project does not hold.
use std::collections::HashSet;
use std::path::{Path, PathBuf};

use crate::loc::language::{LanguageSpec, languages};
use crate::string_mask::starts_in_string_mask;

/// A name an import statement binds in a file.
#[derive(Debug, Clone, PartialEq)]
pub struct Binding {
    /// The module named, with its leading dots.
    pub module: String,
    /// The name `from … import` takes from it, if one is taken.
    pub name: Option<String>,
    /// What the file calls it: the alias, the name taken, or the dotted
    /// path of a plain `import`. `None` for `*`, which binds every name.
    pub bound: Option<String>,
}

impl Binding {
    /// The import as `extract` gives it.
    pub fn import(&self) -> String {
        match &self.name {
            Some(name) => format!("{}:{name}", self.module),
            None => self.module.clone(),
        }
    }

    /// The module `name` would be, were it a submodule of `module`.
    pub fn submodule(&self) -> Option<String> {
        let name = self.name.as_ref()?;
        let dot = if self.module.ends_with('.') { "" } else { "." };
        Some(format!("{}{dot}{name}", self.module))
    }
}

/// A file as its imports and the rest of its code.
#[derive(Debug, Default)]
pub struct Read {
    pub bindings: Vec<Binding>,
    /// The lines that are not import statements, without comments or the
    /// inside of multi-line strings.
    pub code: Vec<String>,
}

/// The specification of Python, which tells where strings are.
pub fn spec() -> Option<&'static LanguageSpec> {
    languages().iter().find(|s| s.name == "Python")
}

/// Extract what each `import` and `from … import` statement names, wherever
/// it is written: at the top, in a function, under `if TYPE_CHECKING:`.
///
/// `import a.b` gives `a.b`. `from a.b import c` gives `a.b:c`, since `c`
/// may be a submodule of `a.b` or a name it defines; `from a.b import *`
/// gives `a.b`. Leading dots are kept: `from . import c` gives `.:c`.
///
/// Statements may run over several lines, in parentheses or with a
/// backslash. A statement written inside a docstring is not one.
pub fn extract(source: &str) -> Vec<String> {
    read(source).bindings.iter().map(Binding::import).collect()
}

/// Read the imports of a file, with the names they bind, apart from the
/// rest of its code.
pub fn read(source: &str) -> Read {
    let mut read = Read::default();
    let Some(spec) = spec() else {
        return read;
    };
    let lines: Vec<String> = source.lines().map(String::from).collect();
    let in_string = starts_in_string_mask(&lines, spec);
    let mut code = lines
        .iter()
        .zip(in_string)
        .filter(|(_, masked)| !masked)
        .map(|(line, _)| without_comment(line));

    while let Some(written) = code.next() {
        let line = written.trim_start();
        if keyword(line, "import").is_none() && keyword(line, "from").is_none() {
            read.code.push(written.to_string());
            continue;
        }
        let mut statement = line.to_string();
        while unfinished(&mut statement) {
            let Some(more) = code.next() else { break };
            statement.push(' ');
            statement.push_str(more);
        }
        for part in statement.split(';') {
            read_statement(part.trim(), &mut read.bindings);
        }
    }
    read
}

/// The line without its comment. An import statement holds no string, so
/// there `#` always starts one.
fn without_comment(line: &str) -> &str {
    line.split('#').next().unwrap_or(line)
}

/// Whether the statement goes on in the next line: it ends in a backslash,
/// which is dropped, or a parenthesis is still open.
fn unfinished(statement: &mut String) -> bool {
    let written = statement.trim_end();
    if written.ends_with('\\') {
        statement.truncate(written.len() - 1);
        return true;
    }
    statement.matches('(').count() > statement.matches(')').count()
}

/// What follows `word` when `text` starts with it as a whole word.
pub fn keyword<'a>(text: &'a str, word: &str) -> Option<&'a str> {
    let rest = text.strip_prefix(word)?;
    let apart = |c: char| !(c.is_alphanumeric() || c == '_');
    rest.starts_with(apart).then_some(rest)
}

/// Add what one statement binds, if it is an import.
fn read_statement(statement: &str, bindings: &mut Vec<Binding>) {
    if let Some(listed) = keyword(statement, "import") {
        bindings.extend(
            names(listed)
                .filter(|(name, _)| is_module(name))
                .map(|(name, alias)| Binding {
                    module: name.to_string(),
                    name: None,
                    bound: Some(alias.unwrap_or(name).to_string()),
                }),
        );
        return;
    }
    let Some(rest) = keyword(statement, "from") else {
        return;
    };
    let Some((module, rest)) = rest.trim_start().split_once(char::is_whitespace) else {
        return;
    };
    let Some(listed) = keyword(rest.trim_start(), "import") else {
        return;
    };
    let path = module.trim_start_matches('.');
    if !(is_module(path) || path.is_empty() && !module.is_empty()) {
        return;
    }
    for (name, alias) in names(listed) {
        let star = name == "*";
        if !star && !is_identifier(name) {
            continue;
        }
        bindings.push(Binding {
            module: module.to_string(),
            name: (!star).then(|| name.to_string()),
            bound: (!star).then(|| alias.unwrap_or(name).to_string()),
        });
    }
}

/// The names of a comma-separated list, each with its `as` alias if it has
/// one, without the parentheses around the list.
fn names(listed: &str) -> impl Iterator<Item = (&str, Option<&str>)> {
    let around = |c: char| c.is_whitespace() || c == '(' || c == ')';
    listed.split(',').filter_map(move |item| {
        let mut words = item.trim_matches(around).split_whitespace();
        let name = words.next()?;
        let alias = words.next().filter(|&word| word == "as").and(words.next());
        Some((name, alias))
    })
}

fn is_identifier(name: &str) -> bool {
    !name.is_empty() && name.chars().all(|c| c.is_alphanumeric() || c == '_')
}

fn is_module(path: &str) -> bool {
    path.split('.').all(is_identifier)
}

/// Resolve what `extract` gives to the project file it names, if any.
///
/// A dotted path names `a/b/c.py`, or the package `a/b/c/__init__.py`; a
/// stub (`.pyi`) stands for a module that has no source. Only the deepest
/// module is the one used: the packages above it are not. For `a.b:c` the
/// submodule `a.b.c` comes first, then `a.b` as the module that defines `c`.
///
/// A relative import starts at the importer's package, one level up per
/// extra dot. An absolute one is looked up from each directory above the
/// importer that is not itself a package, nearest first, and from the `src`
/// directory under each: the roots a project is run or installed from. Then
/// come the roots `declared` by the nearest project above the importer that
/// declares any, each project given as its directory and its roots. A name
/// found under none of them is external.
pub fn resolve(
    importer: &Path,
    import: &str,
    known: &HashSet<PathBuf>,
    declared: &[(PathBuf, Vec<PathBuf>)],
) -> Option<PathBuf> {
    let (module, name) = match import.split_once(':') {
        Some((module, name)) => (module, Some(name)),
        None => (import, None),
    };
    let path = module.trim_start_matches('.');
    let dots = module.len() - path.len();
    let module: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();
    let submodule: Option<Vec<&str>> = name.map(|n| module.iter().copied().chain([n]).collect());

    let dir = importer.parent().unwrap_or(Path::new(""));
    let roots: Vec<PathBuf> = match dots {
        0 => roots(dir, known)
            .into_iter()
            .chain(declared_for(dir, declared).iter().cloned())
            .collect(),
        _ => dir
            .ancestors()
            .nth(dots - 1)
            .map(Path::to_path_buf)
            .into_iter()
            .collect(),
    };
    roots.iter().find_map(|root| {
        submodule
            .iter()
            .chain([&module])
            // Only a relative import can name the package it starts at.
            .filter(|segments| dots > 0 || !segments.is_empty())
            .find_map(|segments| module_file(root, segments, known))
    })
}

/// The directories an absolute import written in `dir` may be looked up
/// from, nearest first.
fn roots(dir: &Path, known: &HashSet<PathBuf>) -> Vec<PathBuf> {
    dir.ancestors()
        .filter(|above| module_file(above, &[], known).is_none())
        .flat_map(|above| [above.to_path_buf(), above.join("src")])
        .collect()
}

/// The roots declared by the nearest project above `dir`.
fn declared_for<'a>(dir: &Path, declared: &'a [(PathBuf, Vec<PathBuf>)]) -> &'a [PathBuf] {
    declared
        .iter()
        .filter(|(project, _)| dir.starts_with(project))
        .max_by_key(|(project, _)| project.components().count())
        .map_or(&[], |(_, roots)| roots)
}

/// The file of the module `segments` name under `root`: a module of its
/// own, or the `__init__` of a package. No segments name `root` itself.
fn module_file(root: &Path, segments: &[&str], known: &HashSet<PathBuf>) -> Option<PathBuf> {
    let package = segments.iter().fold(root.to_path_buf(), |p, s| p.join(s));
    ["py", "pyi"]
        .iter()
        .flat_map(|ext| {
            let file = segments.split_last().map(|(last, above)| {
                let dir = above.iter().fold(root.to_path_buf(), |p, s| p.join(s));
                dir.join(format!("{last}.{ext}"))
            });
            file.into_iter()
                .chain([package.join(format!("__init__.{ext}"))])
        })
        .find(|path| known.contains(path))
}

#[cfg(test)]
#[path = "python_test.rs"]
mod tests;