Skip to main content

oag_fastapi_server/
type_mapper.rs

1use oag_core::ir::IrType;
2
3/// Map an `IrType` to its Python type string representation.
4pub fn ir_type_to_python(ir_type: &IrType) -> String {
5    match ir_type {
6        IrType::String => "str".to_string(),
7        IrType::Number => "float".to_string(),
8        IrType::Integer => "int".to_string(),
9        IrType::Boolean => "bool".to_string(),
10        IrType::Null => "None".to_string(),
11        IrType::DateTime => "str".to_string(),
12        IrType::Binary => "bytes".to_string(),
13        IrType::Any => "Any".to_string(),
14        IrType::Void => "None".to_string(),
15        IrType::Ref(name) => name.clone(),
16        IrType::Array(inner) => {
17            let inner_py = ir_type_to_python(inner);
18            format!("list[{inner_py}]")
19        }
20        IrType::Map(value_type) => {
21            let value_py = ir_type_to_python(value_type);
22            format!("dict[str, {value_py}]")
23        }
24        IrType::Object(fields) => {
25            if fields.is_empty() {
26                return "dict[str, Any]".to_string();
27            }
28            // Inline objects become dict[str, Any] in Python
29            "dict[str, Any]".to_string()
30        }
31        IrType::Union(variants) => {
32            let variant_strs: Vec<String> = variants.iter().map(ir_type_to_python).collect();
33            variant_strs.join(" | ")
34        }
35    }
36}
37
38/// Map an `IrType` to a Python type that's Optional if not required.
39pub fn ir_type_to_python_field(ir_type: &IrType, required: bool) -> String {
40    let base = ir_type_to_python(ir_type);
41    if required {
42        base
43    } else {
44        format!("{base} | None = None")
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test_primitives() {
54        assert_eq!(ir_type_to_python(&IrType::String), "str");
55        assert_eq!(ir_type_to_python(&IrType::Number), "float");
56        assert_eq!(ir_type_to_python(&IrType::Integer), "int");
57        assert_eq!(ir_type_to_python(&IrType::Boolean), "bool");
58        assert_eq!(ir_type_to_python(&IrType::Null), "None");
59        assert_eq!(ir_type_to_python(&IrType::Any), "Any");
60        assert_eq!(ir_type_to_python(&IrType::Void), "None");
61    }
62
63    #[test]
64    fn test_array() {
65        assert_eq!(
66            ir_type_to_python(&IrType::Array(Box::new(IrType::String))),
67            "list[str]"
68        );
69    }
70
71    #[test]
72    fn test_map() {
73        assert_eq!(
74            ir_type_to_python(&IrType::Map(Box::new(IrType::String))),
75            "dict[str, str]"
76        );
77    }
78
79    #[test]
80    fn test_ref() {
81        assert_eq!(ir_type_to_python(&IrType::Ref("Pet".to_string())), "Pet");
82    }
83
84    #[test]
85    fn test_union() {
86        assert_eq!(
87            ir_type_to_python(&IrType::Union(vec![IrType::String, IrType::Integer])),
88            "str | int"
89        );
90    }
91
92    #[test]
93    fn test_optional_field() {
94        assert_eq!(ir_type_to_python_field(&IrType::String, true), "str");
95        assert_eq!(
96            ir_type_to_python_field(&IrType::String, false),
97            "str | None = None"
98        );
99    }
100}