feat: Add Pydantic BaseModel and dataclass support for Python type inference

backend/parsers/windmill-parser-py/src/lib.rs

@@ -15,13 +15,17 @@ use windmill_parser::{json_to_typ, Arg, MainArgSignature, ObjectType, Typ};
 
 use rustpython_parser::{
     ast::{
-        Constant, Expr, ExprConstant, ExprDict, ExprList, ExprName, Stmt, StmtFunctionDef, Suite,
+        Constant, Expr, ExprConstant, ExprDict, ExprList, ExprName,
+        Stmt, StmtFunctionDef, Suite,
     },
     Parse,
 };
 
 pub mod asset_parser;
+pub mod pydantic_parser;
+
 pub use asset_parser::parse_assets;
+use pydantic_parser::ModuleGuard;
 
 const FUNCTION_CALL: &str = "<function call>";
 
@@ -71,6 +75,16 @@ pub fn parse_python_signature(
 
     let has_preprocessor = !filter_non_main(code, "preprocessor").is_empty();
 
+    // Parse full code to get all class definitions for Pydantic/dataclass detection
+    // Use RAII guard to ensure cleanup even on early return
+    let _guard = Suite::parse(code, "main.py")
+        .map_err(|e| {
+            eprintln!("Warning: Failed to parse code for Pydantic/dataclass detection: {:?}", e);
+            e
+        })
+        .ok()
+        .map(ModuleGuard::new);
+
     let filtered_code = filter_non_main(code, &main_name);
     if filtered_code.is_empty() {
         return Ok(MainArgSignature {
@@ -161,6 +175,7 @@ pub fn parse_python_signature(
             has_preprocessor: Some(has_preprocessor),
         })
     }
+    // ModuleGuard automatically calls clear_current_module() when _guard goes out of scope
 }
 
 fn parse_expr(e: &Box<Expr>) -> (Typ, bool) {
@@ -238,7 +253,17 @@ fn parse_typ(id: &str) -> Typ {
         x @ _ if x.starts_with("DynMultiselect_") => {
             Typ::DynMultiselect(x.strip_prefix("DynMultiselect_").unwrap().to_string())
         }
-        _ => Typ::Resource(map_resource_name(id)),
+        _ => {
+            // Check if it's a Pydantic model or dataclass
+            if let Some(module) = pydantic_parser::get_current_module() {
+                if let Some(object_type) = pydantic_parser::detect_model_type(id, &module) {
+                    return Typ::Object(object_type);
+                }
+            }
+
+            // Fallback to Resource if not a model
+            Typ::Resource(map_resource_name(id))
+        }
     }
 }
 
@@ -249,6 +274,7 @@ fn map_resource_name(x: &str) -> String {
     }
 }
 
+
 fn to_value<R>(et: &Expr<R>) -> Option<serde_json::Value> {
     match et {
         Expr::Constant(ExprConstant { value, .. }) => Some(constant_to_value(value)),