use typewriter_core::ir::*;
use typewriter_core::mapper::TypeMapper;
use typewriter_core::naming::{to_file_style, FileStyle};
use crate::emitter;
pub struct PythonMapper {
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.1. 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\""));
}
}