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        IrType::Intersection(parts) => {
37            // Python doesn't have a native intersection type; use the first part as a fallback
38            if parts.len() == 1 {
39                ir_type_to_python(&parts[0])
40            } else {
41                // Multiple inheritance: tuple of base classes
42                let part_strs: Vec<String> = parts.iter().map(ir_type_to_python).collect();
43                part_strs.join(", ")
44            }
45        }
46    }
47}
48
49/// Map an `IrType` to a Python type that's Optional if not required.
50pub fn ir_type_to_python_field(ir_type: &IrType, required: bool) -> String {
51    let base = ir_type_to_python(ir_type);
52    if required {
53        base
54    } else {
55        format!("{base} | None = None")
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62
63    #[test]
64    fn test_primitives() {
65        assert_eq!(ir_type_to_python(&IrType::String), "str");
66        assert_eq!(ir_type_to_python(&IrType::Number), "float");
67        assert_eq!(ir_type_to_python(&IrType::Integer), "int");
68        assert_eq!(ir_type_to_python(&IrType::Boolean), "bool");
69        assert_eq!(ir_type_to_python(&IrType::Null), "None");
70        assert_eq!(ir_type_to_python(&IrType::Any), "Any");
71        assert_eq!(ir_type_to_python(&IrType::Void), "None");
72    }
73
74    #[test]
75    fn test_array() {
76        assert_eq!(
77            ir_type_to_python(&IrType::Array(Box::new(IrType::String))),
78            "list[str]"
79        );
80    }
81
82    #[test]
83    fn test_map() {
84        assert_eq!(
85            ir_type_to_python(&IrType::Map(Box::new(IrType::String))),
86            "dict[str, str]"
87        );
88    }
89
90    #[test]
91    fn test_ref() {
92        assert_eq!(ir_type_to_python(&IrType::Ref("Pet".to_string())), "Pet");
93    }
94
95    #[test]
96    fn test_union() {
97        assert_eq!(
98            ir_type_to_python(&IrType::Union(vec![IrType::String, IrType::Integer])),
99            "str | int"
100        );
101    }
102
103    #[test]
104    fn test_optional_field() {
105        assert_eq!(ir_type_to_python_field(&IrType::String, true), "str");
106        assert_eq!(
107            ir_type_to_python_field(&IrType::String, false),
108            "str | None = None"
109        );
110    }
111}