normalize_languages/

python.rs

//! Python language support.

use crate::{ContainerBody, Import, Language, LanguageSymbols, Visibility};
use tree_sitter::Node;

// ============================================================================
// Python language support
// ============================================================================

/// Python language support.
pub struct Python;

impl Language for Python {
    fn name(&self) -> &'static str {
        "Python"
    }
    fn extensions(&self) -> &'static [&'static str] {
        &["py", "pyi", "pyw"]
    }
    fn grammar_name(&self) -> &'static str {
        "python"
    }

    fn as_symbols(&self) -> Option<&dyn LanguageSymbols> {
        Some(self)
    }

    fn extract_docstring(&self, node: &Node, content: &str) -> Option<String> {
        extract_docstring(node, content)
    }

    fn extract_implements(&self, node: &Node, content: &str) -> crate::ImplementsInfo {
        let mut implements = Vec::new();
        if let Some(superclasses) = node.child_by_field_name("superclasses") {
            let mut cursor = superclasses.walk();
            for child in superclasses.children(&mut cursor) {
                if child.kind() == "identifier" {
                    implements.push(content[child.byte_range()].to_string());
                }
            }
        }
        crate::ImplementsInfo {
            is_interface: false,
            implements,
        }
    }

    fn build_signature(&self, node: &Node, content: &str) -> String {
        let name = match self.node_name(node, content) {
            Some(n) => n,
            None => {
                return content[node.byte_range()]
                    .lines()
                    .next()
                    .unwrap_or("")
                    .trim()
                    .to_string();
            }
        };

        if node.kind() == "class_definition" {
            let bases = node
                .child_by_field_name("superclasses")
                .map(|b| &content[b.byte_range()])
                .unwrap_or("");
            if bases.is_empty() {
                format!("class {}", name)
            } else {
                format!("class {}{}", name, bases)
            }
        } else {
            // function_definition / decorated_definition
            let is_async = node
                .child(0)
                .map(|c| &content[c.byte_range()] == "async")
                .unwrap_or(false);
            let prefix = if is_async { "async def" } else { "def" };
            let params = node
                .child_by_field_name("parameters")
                .map(|p| &content[p.byte_range()])
                .unwrap_or("()");
            let return_type = node
                .child_by_field_name("return_type")
                .map(|r| format!(" -> {}", &content[r.byte_range()]))
                .unwrap_or_default();
            format!("{} {}{}{}", prefix, name, params, return_type)
        }
    }

    fn extract_imports(&self, node: &Node, content: &str) -> Vec<Import> {
        let line = node.start_position().row + 1;

        match node.kind() {
            "import_statement" => {
                // import foo, import foo as bar
                let mut imports = Vec::new();
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    if child.kind() == "dotted_name" {
                        let module = content[child.byte_range()].to_string();
                        imports.push(Import {
                            module,
                            names: Vec::new(),
                            alias: None,
                            is_wildcard: false,
                            is_relative: false,
                            line,
                        });
                    } else if child.kind() == "aliased_import"
                        && let Some(name) = child.child_by_field_name("name")
                    {
                        let module = content[name.byte_range()].to_string();
                        let alias = child
                            .child_by_field_name("alias")
                            .map(|a| content[a.byte_range()].to_string());
                        imports.push(Import {
                            module,
                            names: Vec::new(),
                            alias,
                            is_wildcard: false,
                            is_relative: false,
                            line,
                        });
                    }
                }
                imports
            }
            "import_from_statement" => {
                // from foo import bar, baz
                let module = node
                    .child_by_field_name("module_name")
                    .map(|m| content[m.byte_range()].to_string())
                    .unwrap_or_default();

                // Check for relative import (from . or from .. or from .foo)
                let text = &content[node.byte_range()];
                let is_relative = text.starts_with("from .");

                let mut names = Vec::new();
                let mut is_wildcard = false;
                let module_end = node
                    .child_by_field_name("module_name")
                    .map(|m| m.end_byte())
                    .unwrap_or(0);

                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    match child.kind() {
                        "dotted_name" | "identifier" if child.start_byte() > module_end => {
                            names.push(content[child.byte_range()].to_string());
                        }
                        "aliased_import" => {
                            if let Some(name) = child.child_by_field_name("name") {
                                names.push(content[name.byte_range()].to_string());
                            }
                        }
                        "wildcard_import" => {
                            is_wildcard = true;
                        }
                        _ => {}
                    }
                }

                vec![Import {
                    module,
                    names,
                    alias: None,
                    is_wildcard,
                    is_relative,
                    line,
                }]
            }
            _ => Vec::new(),
        }
    }

    fn format_import(&self, import: &Import, names: Option<&[&str]>) -> String {
        let names_to_use: Vec<&str> = names
            .map(|n| n.to_vec())
            .unwrap_or_else(|| import.names.iter().map(|s| s.as_str()).collect());

        if import.is_wildcard {
            format!("from {} import *", import.module)
        } else if names_to_use.is_empty() {
            if let Some(ref alias) = import.alias {
                format!("import {} as {}", import.module, alias)
            } else {
                format!("import {}", import.module)
            }
        } else {
            format!("from {} import {}", import.module, names_to_use.join(", "))
        }
    }

    fn extract_attributes(&self, node: &Node, content: &str) -> Vec<String> {
        extract_decorators(node, content)
    }

    fn get_visibility(&self, node: &Node, content: &str) -> Visibility {
        if let Some(name) = self.node_name(node, content) {
            if name.starts_with("__") && name.ends_with("__") {
                Visibility::Public // dunder methods
            } else if name.starts_with("__") {
                Visibility::Private // name mangled
            } else if name.starts_with('_') {
                Visibility::Protected // convention private
            } else {
                Visibility::Public
            }
        } else {
            Visibility::Public
        }
    }

    fn is_test_symbol(&self, symbol: &crate::Symbol) -> bool {
        let name = symbol.name.as_str();
        match symbol.kind {
            crate::SymbolKind::Function | crate::SymbolKind::Method => name.starts_with("test_"),
            crate::SymbolKind::Class => name.starts_with("Test") && name.len() > 4,
            crate::SymbolKind::Module => name == "tests" || name == "test" || name == "__tests__",
            _ => false,
        }
    }

    fn test_file_globs(&self) -> &'static [&'static str] {
        &["**/test_*.py", "**/*_test.py"]
    }

    fn extract_module_doc(&self, src: &str) -> Option<String> {
        extract_python_module_doc(src)
    }

    fn body_has_docstring(&self, body: &Node, content: &str) -> bool {
        let _ = content;
        body.child(0)
            .map(|c| {
                c.kind() == "string"
                    || (c.kind() == "expression_statement"
                        && c.child(0).map(|n| n.kind() == "string").unwrap_or(false))
            })
            .unwrap_or(false)
    }

    fn container_body<'a>(&self, node: &'a Node<'a>) -> Option<Node<'a>> {
        node.child_by_field_name("body")
    }

    fn analyze_container_body(
        &self,
        body_node: &Node,
        content: &str,
        inner_indent: &str,
    ) -> Option<ContainerBody> {
        let mut cursor = body_node.walk();
        let children: Vec<_> = body_node.children(&mut cursor).collect();

        if children.is_empty() {
            return Some(ContainerBody {
                content_start: body_node.start_byte(),
                content_end: body_node.end_byte(),
                inner_indent: inner_indent.to_string(),
                is_empty: true,
            });
        }

        let mut first_real_idx = 0;
        for (i, child) in children.iter().enumerate() {
            let is_docstring = if child.kind() == "expression_statement" {
                let mut child_cursor = child.walk();
                child
                    .children(&mut child_cursor)
                    .next()
                    .map(|fc| fc.kind() == "string")
                    .unwrap_or(false)
            } else {
                child.kind() == "string"
            };
            if is_docstring && i == 0 {
                first_real_idx = i + 1;
                continue;
            }
            break;
        }

        let is_empty = children.iter().skip(first_real_idx).all(|c| {
            c.kind() == "pass_statement"
                || c.kind() == "string"
                || (c.kind() == "expression_statement"
                    && c.child(0).map(|fc| fc.kind() == "string").unwrap_or(false))
        });

        let content_start = if first_real_idx < children.len() {
            let child_start = children[first_real_idx].start_byte();
            content[..child_start]
                .rfind('\n')
                .map(|i| i + 1)
                .unwrap_or(child_start)
        } else if !children.is_empty() {
            // normalize-syntax-allow: rust/unwrap-in-impl - !children.is_empty() guarantees last() is Some
            let last_end = children.last().unwrap().end_byte();
            if last_end < content.len() && content.as_bytes()[last_end] == b'\n' {
                last_end + 1
            } else {
                last_end
            }
        } else {
            body_node.start_byte()
        };

        Some(ContainerBody {
            content_start,
            content_end: body_node.end_byte(),
            inner_indent: inner_indent.to_string(),
            is_empty,
        })
    }
}

impl LanguageSymbols for Python {}

/// Extract the module-level docstring from Python source.
///
/// Skips shebang lines and coding-declaration comments, then looks for a
/// triple-quoted string as the first non-comment, non-blank content.
fn extract_python_module_doc(src: &str) -> Option<String> {
    let mut lines = src.lines().peekable();
    // Skip shebang and coding comments (PEP 263)
    loop {
        match lines.peek() {
            Some(line) => {
                let t = line.trim();
                if t.starts_with("#!") || t.starts_with("# -*-") || t.starts_with("# coding") {
                    lines.next();
                } else {
                    break;
                }
            }
            None => return None,
        }
    }
    let remaining: String = lines.collect::<Vec<_>>().join("\n");
    let trimmed = remaining.trim_start();

    // Must start with triple-quote string
    let (quote, rest) = if let Some(rest) = trimmed.strip_prefix("\"\"\"") {
        ("\"\"\"", rest)
    } else if let Some(rest) = trimmed.strip_prefix("'''") {
        ("'''", rest)
    } else {
        return None;
    };

    // Find the closing triple-quote
    let end = rest.find(quote)?;
    let doc = rest[..end].trim();
    if doc.is_empty() {
        None
    } else {
        Some(doc.to_string())
    }
}

/// Extract a Python docstring from a function or class body.
///
/// Looks for the first statement in the body being a string literal.
/// Handles both old grammar style (expression_statement > string) and
/// new arborium style (string directly, with string_content child).
fn extract_docstring(node: &Node, content: &str) -> Option<String> {
    let body = node.child_by_field_name("body")?;
    let first = body.child(0)?;

    // Handle both grammar versions:
    // - Old: expression_statement > string
    // - New (arborium): string directly, with string_content child
    let string_node = match first.kind() {
        "string" => Some(first),
        "expression_statement" => first.child(0).filter(|n| n.kind() == "string"),
        _ => None,
    }?;

    // Try string_content child (arborium style)
    let mut cursor = string_node.walk();
    for child in string_node.children(&mut cursor) {
        if child.kind() == "string_content" {
            let doc = content[child.byte_range()].trim();
            if !doc.is_empty() {
                return Some(doc.to_string());
            }
        }
    }

    // Fallback: extract from full string text (old style)
    let text = &content[string_node.byte_range()];
    let doc = text
        .trim_start_matches("\"\"\"")
        .trim_start_matches("'''")
        .trim_start_matches('"')
        .trim_start_matches('\'')
        .trim_end_matches("\"\"\"")
        .trim_end_matches("'''")
        .trim_end_matches('"')
        .trim_end_matches('\'')
        .trim();

    if !doc.is_empty() {
        Some(doc.to_string())
    } else {
        None
    }
}

/// Extract decorators from a Python definition node.
/// Python wraps decorated definitions in a `decorated_definition` parent node.
/// The node passed here is `function_definition` or `class_definition`,
/// so we look at the parent for `decorator` siblings.
fn extract_decorators(node: &Node, content: &str) -> Vec<String> {
    let mut attrs = Vec::new();
    if let Some(parent) = node.parent()
        && parent.kind() == "decorated_definition"
    {
        let mut cursor = parent.walk();
        for child in parent.children(&mut cursor) {
            if child.kind() == "decorator" {
                attrs.push(content[child.byte_range()].to_string());
            }
        }
    }
    attrs
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::GrammarLoader;
    use tree_sitter::Parser;

    struct ParseResult {
        tree: tree_sitter::Tree,
        #[allow(dead_code)]
        loader: GrammarLoader,
    }

    fn parse_python(content: &str) -> ParseResult {
        let loader = GrammarLoader::new();
        let language = loader.get("python").ok().unwrap();
        let mut parser = Parser::new();
        parser.set_language(&language).unwrap();
        ParseResult {
            tree: parser.parse(content, None).unwrap(),
            loader,
        }
    }

    #[test]
    fn test_python_extract_function() {
        let support = Python;
        let content = r#"def foo(x: int) -> str:
    """Convert to string."""
    return str(x)
"#;
        let result = parse_python(content);
        let root = result.tree.root_node();

        // Find function node
        let mut cursor = root.walk();
        let func = root
            .children(&mut cursor)
            .find(|n| n.kind() == "function_definition")
            .unwrap();

        let sig = support.build_signature(&func, content);
        let doc = support.extract_docstring(&func, content);
        assert_eq!(support.node_name(&func, content), Some("foo"));
        assert!(sig.contains("def foo(x: int) -> str"));
        assert_eq!(doc, Some("Convert to string.".to_string()));
    }

    #[test]
    fn test_python_extract_class() {
        let support = Python;
        let content = r#"class Foo(Bar):
    """A foo class."""
    pass
"#;
        let result = parse_python(content);
        let root = result.tree.root_node();

        let mut cursor = root.walk();
        let class = root
            .children(&mut cursor)
            .find(|n| n.kind() == "class_definition")
            .unwrap();

        let sig = support.build_signature(&class, content);
        let doc = support.extract_docstring(&class, content);
        assert_eq!(support.node_name(&class, content), Some("Foo"));
        assert!(sig.contains("class Foo(Bar)"));
        assert_eq!(doc, Some("A foo class.".to_string()));
    }

    #[test]
    fn test_python_visibility() {
        let support = Python;
        let content = r#"def public(): pass
def _protected(): pass
def __private(): pass
def __dunder__(): pass
"#;
        let result = parse_python(content);
        let root = result.tree.root_node();

        let mut cursor = root.walk();
        let funcs: Vec<_> = root
            .children(&mut cursor)
            .filter(|n| n.kind() == "function_definition")
            .collect();

        assert_eq!(
            support.get_visibility(&funcs[0], content),
            Visibility::Public
        );
        assert_eq!(
            support.get_visibility(&funcs[1], content),
            Visibility::Protected
        );
        assert_eq!(
            support.get_visibility(&funcs[2], content),
            Visibility::Private
        );
        assert_eq!(
            support.get_visibility(&funcs[3], content),
            Visibility::Public
        ); // dunder
    }

    /// Documents node kinds that exist in the Python grammar but aren't used in trait methods.
    /// Each exclusion has a reason. Review periodically as features expand.
    ///
    /// Run `cross_check_node_kinds` in registry.rs to see all potentially useful kinds.
    #[test]
    fn unused_node_kinds_audit() {
        use crate::validate_unused_kinds_audit;

        // Categories:
        // - STRUCTURAL: Internal/wrapper nodes, not semantically meaningful on their own
        // - CLAUSE: Sub-parts of statements, handled via parent (e.g., else_clause in if_statement)
        // - EXPRESSION: Expressions don't create control flow/scope, we track statements
        // - TYPE: Type annotation nodes, not relevant for current analysis
        // - LEGACY: Python 2 compatibility, not worth supporting
        // - OPERATOR: Operators within expressions, too granular
        // - MAYBE: Potentially useful, to be added when needed

        #[rustfmt::skip]
        let documented_unused: &[&str] = &[
            // STRUCTURAL
            "aliased_import",          // used internally by extract_imports
            "block",                   // generic block wrapper (duplicate in grammar)
            "expression_list",         // comma-separated expressions              // too common, used everywhere
            "import_prefix",           // dots in relative imports
            "lambda_parameters",       // internal to lambda                  // root node of file
            "parenthesized_expression",// grouping only
            "relative_import",         // handled in extract_imports
            "tuple_expression",        // comma-separated values
            "wildcard_import",         // handled in extract_imports

            // CLAUSE (sub-parts of statements)
            "case_pattern",            // internal to case_clause
            "class_pattern",           // pattern in match/case
            "elif_clause",             // part of if_statement
            "else_clause",             // part of if/for/while/try
            "finally_clause",          // part of try_statement
            "for_in_clause",           // internal to comprehensions
            "if_clause",               // internal to comprehensions
            "with_clause",             // internal to with_statement
            "with_item",               // internal to with_statement

            // EXPRESSION (don't affect control flow structure)
            "await",                   // await keyword, not a statement
            "format_expression",       // f-string interpolation
            "format_specifier",        // f-string format spec
            "named_expression",        // walrus operator :=
            "yield",                   // yield keyword form

            // TYPE (type annotations)
            "constrained_type",        // type constraints
            "generic_type",            // parameterized types
            "member_type",             // attribute access in types
            "splat_type",              // *args/**kwargs types
            "type",                    // generic type node
            "type_alias_statement",    // could track as symbol
            "type_conversion",         // !r/!s/!a in f-strings
            "type_parameter",          // generic type params
            "typed_default_parameter", // param with type and default
            "typed_parameter",         // param with type annotation
            "union_type",              // X | Y union syntax

            // OPERATOR
            "binary_operator",         // +, -, *, /, etc.
            "boolean_operator",        // and/or - handled in complexity_nodes as keywords
            "comparison_operator",     // ==, <, >, etc.
            "not_operator",            // not keyword
            "unary_operator",          // -, +, ~

            // LEGACY (Python 2)
            "exec_statement",          // Python 2 exec
            "print_statement",         // Python 2 print

            // MAYBE: Potentially useful
            "decorated_definition",    // wrapper for @decorator
            "delete_statement",        // del statement
            "future_import_statement", // from __future__
            "global_statement",        // scope modifier
            "nonlocal_statement",      // scope modifier
            "pass_statement",          // no-op, detect empty bodies
            // control flow — not extracted as symbols
            "lambda",
            "import_statement",
            "continue_statement",
            "raise_statement",
            "case_clause",
            "generator_expression",
            "assert_statement",
            "if_statement",
            "while_statement",
            "with_statement",
            "try_statement",
            "import_from_statement",
            "return_statement",
            "except_clause",
            "dictionary_comprehension",
            "conditional_expression",
            "match_statement",
            "set_comprehension",
            "for_statement",
            "list_comprehension",
            "break_statement",
        ];

        validate_unused_kinds_audit(&Python, documented_unused)
            .expect("Python unused node kinds audit failed");
    }
}