pleme-doc-gen 0.1.40

//! Typed Python AST + pretty-printer (config-class subset).
//!
//! Per the ★★★ NO-format!()-for-code prime directive. Covers the
//! Python-class shape used in `conanfile.py` and similar config-as-
//! class DSLs (setup.py, noxfile.py, fabfile.py):
//!
//!   from conan import ConanFile
//!
//!   class DemoConan(ConanFile):
//!       name = "demo"
//!       version = "0.1.0"
//!       settings = "os", "compiler", "build_type", "arch"
//!
//! NOT a general Python AST — no control flow, no function bodies
//! beyond bare statements. Future need: extend by adding variants,
//! never format!().

use std::fmt::Write;

#[derive(Debug, Clone)]
pub enum Expr {
    /// `"text"` — double-quoted with Python's escape rules.
    String(String),
    Int(i64),
    Bool(bool),
    None,
    /// Bare identifier (variable / class / type name).
    Ident(String),
    /// `[a, b, c]` — list.
    List(Vec<Expr>),
    /// `(a, b, c)` — tuple. Single-element tuples get a trailing comma.
    Tuple(Vec<Expr>),
    /// `{k: v, k: v}` — dict.
    Dict(Vec<(Expr, Expr)>),
    /// `func(arg, kw=val)` — call with positional + keyword args.
    Call { name: String, args: Vec<(Option<String>, Expr)> },
}

impl Expr {
    pub fn s(v: impl Into<String>) -> Self { Self::String(v.into()) }
    pub fn i(v: i64) -> Self { Self::Int(v) }
    pub fn b(v: bool) -> Self { Self::Bool(v) }
    pub fn ident(v: impl Into<String>) -> Self { Self::Ident(v.into()) }
    pub fn list(items: impl IntoIterator<Item = Expr>) -> Self {
        Self::List(items.into_iter().collect())
    }
    pub fn tuple(items: impl IntoIterator<Item = Expr>) -> Self {
        Self::Tuple(items.into_iter().collect())
    }
    pub fn call(name: impl Into<String>, args: Vec<(Option<String>, Expr)>) -> Self {
        Self::Call { name: name.into(), args }
    }
    pub fn pos(e: Expr) -> (Option<String>, Expr) { (None, e) }
    pub fn kw(name: &str, e: Expr) -> (Option<String>, Expr) {
        (Some(name.to_string()), e)
    }

    fn render(&self) -> String {
        match self {
            Self::String(s) => render_python_string(s),
            Self::Int(n) => n.to_string(),
            Self::Bool(true) => "True".into(),
            Self::Bool(false) => "False".into(),
            Self::None => "None".into(),
            Self::Ident(s) => s.clone(),
            Self::List(items) => {
                let parts: Vec<String> = items.iter().map(Expr::render).collect();
                let mut s = String::from("[");
                s.push_str(&parts.join(", "));
                s.push(']');
                s
            }
            Self::Tuple(items) => {
                let parts: Vec<String> = items.iter().map(Expr::render).collect();
                // Python single-element tuple: `(a,)`; >=2: `(a, b)`.
                // For setup-style `name = "a", "b", "c"` callers should
                // use Tuple too — the surrounding context drops the parens
                // when it's a top-level assign RHS (handled in Stmt::Assign).
                let joined = parts.join(", ");
                if items.len() == 1 {
                    let mut s = String::from("("); s.push_str(&joined); s.push_str(",)"); s
                } else {
                    let mut s = String::from("("); s.push_str(&joined); s.push(')'); s
                }
            }
            Self::Dict(items) => {
                let parts: Vec<String> = items.iter().map(|(k, v)| {
                    let mut s = k.render(); s.push_str(": "); s.push_str(&v.render()); s
                }).collect();
                let mut s = String::from("{");
                s.push_str(&parts.join(", "));
                s.push('}');
                s
            }
            Self::Call { name, args } => {
                let parts: Vec<String> = args.iter().map(|(kw, v)| {
                    let mut p = String::new();
                    if let Some(k) = kw { p.push_str(k); p.push('='); }
                    p.push_str(&v.render());
                    p
                }).collect();
                let mut s = String::from(name); s.push('(');
                s.push_str(&parts.join(", "));
                s.push(')');
                s
            }
        }
    }
}

fn render_python_string(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('"');
    for c in s.chars() {
        match c {
            '\\' => out.push_str(r"\\"),
            '"' => out.push_str("\\\""),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c => out.push(c),
        }
    }
    out.push('"');
    out
}

#[derive(Debug, Clone)]
pub enum Stmt {
    /// `name = value` — module- or class-level assignment.
    /// For tuple RHS the surrounding parens are dropped (Python
    /// "comma-tuple" idiom that `class.settings = "a", "b"` relies on).
    Assign { name: String, value: Expr },
    /// Bare expression statement: `super().__init__()`.
    Expr(Expr),
    /// `from <module> import <name>, <name>` (or just `import <module>`).
    Import { module: String, names: Vec<String> },
    /// `def <name>(<params>):` — function definition with typed body.
    /// Params are positional; type annotations are passed verbatim
    /// in each param string (e.g. "self", "x: int", "y: str = 'a'").
    FunctionDef { name: String, params: Vec<String>, body: Vec<Stmt> },
    /// `if <cond>:` — conditional with typed body. `cond` is rendered
    /// as a raw token (callers compose typed Expr::render output).
    /// Useful for `if __name__ == "__main__":` test entrypoints.
    If { cond: String, body: Vec<Stmt> },
    Blank,
    Comment(String),
}

#[derive(Debug, Clone)]
pub struct Class {
    pub name: String,
    pub bases: Vec<String>,
    pub body: Vec<Stmt>,
}

impl Class {
    pub fn new(name: impl Into<String>, bases: impl IntoIterator<Item = String>) -> Self {
        Self { name: name.into(), bases: bases.into_iter().collect(), body: Vec::new() }
    }
    pub fn push(&mut self, s: Stmt) -> &mut Self { self.body.push(s); self }
}

#[derive(Debug, Clone, Default)]
pub struct File {
    /// Module-level statements that render BEFORE class defs. Use
    /// for `import` lines + module-level constants.
    pub stmts: Vec<Stmt>,
    /// Class definitions render after `stmts`, separated by a blank.
    pub classes: Vec<Class>,
    /// Statements that render AFTER the class defs. Use for
    /// `if __name__ == "__main__":` entrypoints + post-class
    /// invocations that depend on the classes being defined.
    pub post_class_stmts: Vec<Stmt>,
}

impl File {
    pub fn new() -> Self { Self::default() }
    pub fn push(&mut self, s: Stmt) -> &mut Self { self.stmts.push(s); self }
    pub fn class(&mut self, c: Class) -> &mut Self { self.classes.push(c); self }
    pub fn push_after_class(&mut self, s: Stmt) -> &mut Self {
        self.post_class_stmts.push(s); self
    }
    pub fn render(&self) -> String {
        let mut out = String::new();
        render_stmts(&self.stmts, 0, &mut out);
        if !self.classes.is_empty() && !self.stmts.is_empty() {
            out.push('\n');
        }
        for (i, c) in self.classes.iter().enumerate() {
            if i > 0 { out.push('\n'); }
            let bases = if c.bases.is_empty() { String::new() } else {
                let mut s = String::from("(");
                s.push_str(&c.bases.join(", "));
                s.push(')');
                s
            };
            let _ = writeln!(out, "class {}{bases}:", c.name);
            if c.body.is_empty() {
                out.push_str("    pass\n");
            } else {
                render_stmts(&c.body, 1, &mut out);
            }
        }
        if !self.post_class_stmts.is_empty() {
            if !self.classes.is_empty() || !self.stmts.is_empty() {
                out.push('\n');
            }
            render_stmts(&self.post_class_stmts, 0, &mut out);
        }
        out
    }
}

fn render_stmts(stmts: &[Stmt], depth: usize, out: &mut String) {
    let pad = "    ".repeat(depth);
    for s in stmts {
        match s {
            Stmt::Assign { name, value } => {
                // Drop parens for a tuple RHS at assignment context.
                let val = match value {
                    Expr::Tuple(items) if items.len() >= 2 => {
                        let parts: Vec<String> = items.iter().map(Expr::render).collect();
                        parts.join(", ")
                    }
                    _ => value.render(),
                };
                let _ = writeln!(out, "{pad}{name} = {val}");
            }
            Stmt::Expr(e) => { let _ = writeln!(out, "{pad}{}", e.render()); }
            Stmt::FunctionDef { name, params, body } => {
                let _ = writeln!(out, "{pad}def {name}({}):", params.join(", "));
                if body.is_empty() {
                    let _ = writeln!(out, "{}pass", "    ".repeat(depth + 1));
                } else {
                    render_stmts(body, depth + 1, out);
                }
            }
            Stmt::If { cond, body } => {
                let _ = writeln!(out, "{pad}if {cond}:");
                if body.is_empty() {
                    let _ = writeln!(out, "{}pass", "    ".repeat(depth + 1));
                } else {
                    render_stmts(body, depth + 1, out);
                }
            }
            Stmt::Import { module, names } if names.is_empty() => {
                let _ = writeln!(out, "{pad}import {module}");
            }
            Stmt::Import { module, names } => {
                let _ = writeln!(out, "{pad}from {module} import {}", names.join(", "));
            }
            Stmt::Blank => out.push('\n'),
            Stmt::Comment(c) => { let _ = writeln!(out, "{pad}# {c}"); }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn import_renders_from_x_import_y() {
        let mut f = File::new();
        f.push(Stmt::Import {
            module: "conan".into(), names: vec!["ConanFile".into()],
        });
        assert_eq!(f.render(), "from conan import ConanFile\n");
    }
    #[test]
    fn class_with_body_renders_with_4space_indent() {
        let mut f = File::new();
        let mut c = Class::new("Demo", vec!["ConanFile".into()]);
        c.push(Stmt::Assign { name: "name".into(), value: Expr::s("demo") });
        f.class(c);
        let out = f.render();
        assert!(out.contains("class Demo(ConanFile):\n"));
        assert!(out.contains("    name = \"demo\"\n"));
    }
    #[test]
    fn tuple_assign_drops_parens() {
        let mut f = File::new();
        let mut c = Class::new("X", vec![]);
        c.push(Stmt::Assign {
            name: "settings".into(),
            value: Expr::tuple(vec![Expr::s("os"), Expr::s("arch")]),
        });
        f.class(c);
        let out = f.render();
        assert!(out.contains("    settings = \"os\", \"arch\"\n"),
            "tuple-RHS should drop enclosing parens; got: {out}");
    }
    #[test]
    fn function_def_renders_with_typed_body() {
        let mut f = File::new();
        let body = vec![
            Stmt::Expr(Expr::call("self.assertEqual",
                vec![Expr::pos(Expr::call("smoke", vec![])), Expr::pos(Expr::i(4))])),
        ];
        f.push(Stmt::FunctionDef {
            name: "test_smoke".into(),
            params: vec!["self".into()],
            body,
        });
        let out = f.render();
        assert!(out.contains("def test_smoke(self):"));
        assert!(out.contains("    self.assertEqual(smoke(), 4)"));
    }

    #[test]
    fn if_main_guard_renders() {
        let mut f = File::new();
        f.push(Stmt::If {
            cond: r#"__name__ == "__main__""#.into(),
            body: vec![Stmt::Expr(Expr::call("unittest.main", vec![]))],
        });
        let out = f.render();
        assert!(out.contains(r#"if __name__ == "__main__":"#));
        assert!(out.contains("    unittest.main()"));
    }

    #[test]
    fn empty_function_emits_pass() {
        let mut f = File::new();
        f.push(Stmt::FunctionDef {
            name: "noop".into(),
            params: vec![],
            body: vec![],
        });
        let out = f.render();
        assert!(out.contains("def noop():"));
        assert!(out.contains("    pass"));
    }

    #[test]
    fn empty_class_emits_pass() {
        let mut f = File::new();
        f.class(Class::new("Empty", vec![]));
        let out = f.render();
        assert!(out.contains("class Empty:\n    pass\n"));
    }
}