typewriter-python 0.3.0

//! Python type mapper implementation.

use typewriter_core::ir::*;
use typewriter_core::mapper::TypeMapper;
use typewriter_core::naming::{to_file_style, FileStyle};

use crate::emitter;

/// Python language mapper.
///
/// Generates Pydantic v2 `BaseModel` classes from Rust structs and
/// Python `Enum` / `Union` types from Rust enums.
pub struct PythonMapper {
    /// File naming style (default: `snake_case`)
    pub file_style: FileStyle,
}

impl PythonMapper {
    pub fn new() -> Self {
        Self {
            file_style: FileStyle::SnakeCase,
        }
    }

    pub fn with_file_style(mut self, style: FileStyle) -> Self {
        self.file_style = style;
        self
    }
}

impl Default for PythonMapper {
    fn default() -> Self {
        Self::new()
    }
}

impl TypeMapper for PythonMapper {
    fn map_primitive(&self, ty: &PrimitiveType) -> String {
        match ty {
            PrimitiveType::String => "str".to_string(),
            PrimitiveType::Bool => "bool".to_string(),
            PrimitiveType::U8
            | PrimitiveType::U16
            | PrimitiveType::U32
            | PrimitiveType::U64
            | PrimitiveType::U128
            | PrimitiveType::I8
            | PrimitiveType::I16
            | PrimitiveType::I32
            | PrimitiveType::I64
            | PrimitiveType::I128 => "int".to_string(),
            PrimitiveType::F32 | PrimitiveType::F64 => "float".to_string(),
            PrimitiveType::Uuid => "UUID".to_string(),
            PrimitiveType::DateTime => "datetime".to_string(),
            PrimitiveType::NaiveDate => "date".to_string(),
            PrimitiveType::JsonValue => "Any".to_string(),
        }
    }

    fn map_option(&self, inner: &TypeKind) -> String {
        format!("Optional[{}]", self.map_type(inner))
    }

    fn map_vec(&self, inner: &TypeKind) -> String {
        format!("list[{}]", self.map_type(inner))
    }

    fn map_hashmap(&self, key: &TypeKind, value: &TypeKind) -> String {
        format!("dict[{}, {}]", self.map_type(key), self.map_type(value))
    }

    fn map_tuple(&self, elements: &[TypeKind]) -> String {
        let inner: Vec<String> = elements.iter().map(|e| self.map_type(e)).collect();
        format!("tuple[{}]", inner.join(", "))
    }

    fn map_named(&self, name: &str) -> String {
        name.to_string()
    }

    fn emit_struct(&self, def: &StructDef) -> String {
        emitter::render_model(self, def)
    }

    fn emit_enum(&self, def: &EnumDef) -> String {
        emitter::render_enum(self, def)
    }

    fn file_header(&self, type_name: &str) -> String {
        format!(
            "# Auto-generated by typewriter v0.3.0. DO NOT EDIT.\n\
            # Source: {}\n\n",
            type_name
        )
    }

    fn file_extension(&self) -> &str {
        "py"
    }

    fn emit_imports(&self, def: &TypeDef) -> String {
        let refs = def.collect_referenced_types();
        if refs.is_empty() {
            return String::new();
        }
        let mut output = String::new();
        for name in &refs {
            let file_name = self.file_naming(name);
            output.push_str(&format!("from .{} import {}\n", file_name, name));
        }
        output
    }

    fn file_naming(&self, type_name: &str) -> String {
        to_file_style(type_name, self.file_style)
    }

    fn map_generic(&self, name: &str, params: &[TypeKind]) -> String {
        let param_strs: Vec<String> = params.iter().map(|p| self.map_type(p)).collect();
        format!("{}[{}]", name, param_strs.join(", "))
    }

    fn map_type(&self, ty: &TypeKind) -> String {
        match ty {
            TypeKind::Primitive(p) => self.map_primitive(p),
            TypeKind::Option(inner) => self.map_option(inner),
            TypeKind::Vec(inner) => self.map_vec(inner),
            TypeKind::HashMap(k, v) => self.map_hashmap(k, v),
            TypeKind::Tuple(elements) => self.map_tuple(elements),
            TypeKind::Named(name) => self.map_named(name),
            TypeKind::Generic(name, params) => self.map_generic(name, params),
            TypeKind::Unit => "None".to_string(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn mapper() -> PythonMapper {
        PythonMapper::new()
    }

    #[test]
    fn test_primitive_mappings() {
        let m = mapper();
        assert_eq!(m.map_primitive(&PrimitiveType::String), "str");
        assert_eq!(m.map_primitive(&PrimitiveType::Bool), "bool");
        assert_eq!(m.map_primitive(&PrimitiveType::U32), "int");
        assert_eq!(m.map_primitive(&PrimitiveType::I64), "int");
        assert_eq!(m.map_primitive(&PrimitiveType::F64), "float");
        assert_eq!(m.map_primitive(&PrimitiveType::Uuid), "UUID");
        assert_eq!(m.map_primitive(&PrimitiveType::DateTime), "datetime");
        assert_eq!(m.map_primitive(&PrimitiveType::NaiveDate), "date");
        assert_eq!(m.map_primitive(&PrimitiveType::JsonValue), "Any");
    }

    #[test]
    fn test_option_mapping() {
        let m = mapper();
        assert_eq!(
            m.map_option(&TypeKind::Primitive(PrimitiveType::String)),
            "Optional[str]"
        );
    }

    #[test]
    fn test_vec_mapping() {
        let m = mapper();
        assert_eq!(
            m.map_vec(&TypeKind::Primitive(PrimitiveType::U32)),
            "list[int]"
        );
    }

    #[test]
    fn test_hashmap_mapping() {
        let m = mapper();
        assert_eq!(
            m.map_hashmap(
                &TypeKind::Primitive(PrimitiveType::String),
                &TypeKind::Primitive(PrimitiveType::U32)
            ),
            "dict[str, int]"
        );
    }

    #[test]
    fn test_tuple_mapping() {
        let m = mapper();
        assert_eq!(
            m.map_tuple(&[
                TypeKind::Primitive(PrimitiveType::String),
                TypeKind::Primitive(PrimitiveType::Bool)
            ]),
            "tuple[str, bool]"
        );
    }

    #[test]
    fn test_file_naming_snake() {
        let m = mapper();
        assert_eq!(m.file_naming("UserProfile"), "user_profile");
        assert_eq!(m.file_naming("User"), "user");
        assert_eq!(m.file_naming("HTTPResponse"), "http_response");
    }

    #[test]
    fn test_file_naming_kebab() {
        let m = PythonMapper::new().with_file_style(FileStyle::KebabCase);
        assert_eq!(m.file_naming("UserProfile"), "user-profile");
        assert_eq!(m.file_naming("HTTPResponse"), "http-response");
    }

    #[test]
    fn test_file_naming_pascal() {
        let m = PythonMapper::new().with_file_style(FileStyle::PascalCase);
        assert_eq!(m.file_naming("UserProfile"), "UserProfile");
    }

    #[test]
    fn test_output_filename() {
        let m = mapper();
        assert_eq!(m.output_filename("UserProfile"), "user_profile.py");
    }

    #[test]
    fn test_output_filename_pascal() {
        let m = PythonMapper::new().with_file_style(FileStyle::PascalCase);
        assert_eq!(m.output_filename("UserProfile"), "UserProfile.py");
    }

    #[test]
    fn test_emit_simple_struct() {
        let m = mapper();
        let def = StructDef {
            name: "User".to_string(),
            fields: vec![
                FieldDef {
                    name: "id".to_string(),
                    ty: TypeKind::Primitive(PrimitiveType::Uuid),
                    optional: false,
                    rename: None,
                    doc: None,
                    skip: false,
                    flatten: false,
                    type_override: None,
                },
                FieldDef {
                    name: "email".to_string(),
                    ty: TypeKind::Primitive(PrimitiveType::String),
                    optional: false,
                    rename: None,
                    doc: None,
                    skip: false,
                    flatten: false,
                    type_override: None,
                },
                FieldDef {
                    name: "age".to_string(),
                    ty: TypeKind::Option(Box::new(TypeKind::Primitive(PrimitiveType::U32))),
                    optional: true,
                    rename: None,
                    doc: None,
                    skip: false,
                    flatten: false,
                    type_override: None,
                },
            ],
            doc: None,
            generics: vec![],
        };

        let output = m.emit_struct(&def);
        assert!(output.contains("class User(BaseModel):"));
        assert!(output.contains("id: UUID"));
        assert!(output.contains("email: str"));
        assert!(output.contains("age: Optional[int] = None"));
    }

    #[test]
    fn test_skipped_field() {
        let m = mapper();
        let def = StructDef {
            name: "User".to_string(),
            fields: vec![
                FieldDef {
                    name: "email".to_string(),
                    ty: TypeKind::Primitive(PrimitiveType::String),
                    optional: false,
                    rename: None,
                    doc: None,
                    skip: false,
                    flatten: false,
                    type_override: None,
                },
                FieldDef {
                    name: "password_hash".to_string(),
                    ty: TypeKind::Primitive(PrimitiveType::String),
                    optional: false,
                    rename: None,
                    doc: None,
                    skip: true,
                    flatten: false,
                    type_override: None,
                },
            ],
            doc: None,
            generics: vec![],
        };

        let output = m.emit_struct(&def);
        assert!(output.contains("email: str"));
        assert!(!output.contains("password_hash"));
    }

    #[test]
    fn test_simple_enum() {
        let m = mapper();
        let def = EnumDef {
            name: "Role".to_string(),
            variants: vec![
                VariantDef {
                    name: "Admin".to_string(),
                    rename: None,
                    kind: VariantKind::Unit,
                    doc: None,
                },
                VariantDef {
                    name: "User".to_string(),
                    rename: None,
                    kind: VariantKind::Unit,
                    doc: None,
                },
            ],
            representation: EnumRepr::External,
            doc: None,
        };

        let output = m.emit_enum(&def);
        assert!(output.contains("class Role(str, Enum):"));
        assert!(output.contains("ADMIN = \"Admin\""));
        assert!(output.contains("USER = \"User\""));
    }
}