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