diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index 4c45b743e5..159b54ffb4 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/backend/parsers/windmill-parser-py/src/lib.rs @@ -15,16 +15,31 @@ use windmill_parser::{json_to_typ, Arg, MainArgSignature, ObjectType, Typ}; use rustpython_parser::{ ast::{ - Constant, Expr, ExprAttribute, ExprConstant, ExprDict, ExprList, ExprName, Stmt, StmtAssign, StmtClassDef, StmtFunctionDef, Suite, + Constant, Expr, ExprAttribute, ExprConstant, ExprDict, ExprList, ExprName, Stmt, + StmtAssign, StmtClassDef, StmtFunctionDef, Suite, }, Parse, }; pub mod asset_parser; +pub mod pydantic_parser; + pub use asset_parser::parse_assets; const FUNCTION_CALL: &str = ""; +/// Cheap string-based check to see if code might contain Pydantic models or dataclasses. +/// Returns true if we should do full AST parsing for type detection, false otherwise. +/// This avoids expensive parsing for the common case where scripts don't use these features. +fn should_parse_for_models(code: &str) -> bool { + code.contains("BaseModel") + || code.contains("from pydantic") + || code.contains("import pydantic") + || code.contains("@dataclass") + || code.contains("from dataclasses") + || code.contains("import dataclasses") +} + fn filter_non_main(code: &str, main_name: &str) -> String { let def_main = format!("def {}(", main_name); let mut filtered_code = String::new(); @@ -139,9 +154,15 @@ fn extract_code_metadata(code: &str, main_name: &str) -> CodeMetadata { for item in body { if let Stmt::Assign(StmtAssign { targets, value, .. }) = item { - if let Some(Expr::Name(ExprName { id: target_name, .. })) = targets.first() { + if let Some(Expr::Name(ExprName { id: target_name, .. })) = + targets.first() + { if !target_name.starts_with('_') { - if let Expr::Constant(ExprConstant { value: Constant::Str(val), .. }) = value.as_ref() { + if let Expr::Constant(ExprConstant { + value: Constant::Str(val), + .. + }) = value.as_ref() + { values.push(val.to_string()); members.insert(target_name.to_string(), val.to_string()); } @@ -154,16 +175,20 @@ fn extract_code_metadata(code: &str, main_name: &str) -> CodeMetadata { enums.insert(name.to_string(), EnumInfo { values, members }); } } - }, + } Stmt::FunctionDef(StmtFunctionDef { name: func_name, body, .. }) if has_docstring => { if &func_name == main_name { if let Some(Stmt::Expr(expr_stmt)) = body.first() { - if let Expr::Constant(ExprConstant { value: Constant::Str(docstring), .. }) = expr_stmt.value.as_ref() { + if let Expr::Constant(ExprConstant { + value: Constant::Str(docstring), + .. + }) = expr_stmt.value.as_ref() + { descriptions = parse_docstring_args(docstring); } } } - }, + } _ => {} } } @@ -231,8 +256,46 @@ pub fn parse_python_signature( let has_preprocessor = !filter_non_main(code, "preprocessor").is_empty(); - let filtered_code = filter_non_main(code, &main_name); - if filtered_code.is_empty() { + // Optimization: Parse code only once + // - If models detected: parse full code, extract main from it, keep AST for type detection + // - If no models: parse only the filtered main function + let (params, module) = if should_parse_for_models(code) { + // Parse full code once for both Pydantic detection and signature extraction + let ast = Suite::parse(code, "main.py") + .map_err(|e| anyhow::anyhow!("Error parsing code: {}", e.to_string()))?; + + // Extract main function from full AST + let params = ast.iter().find_map(|x| match x { + Stmt::FunctionDef(StmtFunctionDef { name, args, .. }) if name == &main_name => { + Some(args.as_ref().clone()) + } + _ => None, + }); + + // Keep AST for Pydantic/dataclass detection + (params, Some(ast)) + } else { + // No models detected - parse only the filtered main function + let filtered_code = filter_non_main(code, &main_name); + if filtered_code.is_empty() { + (None, None) + } else { + let ast = Suite::parse(&filtered_code, "main.py") + .map_err(|e| anyhow::anyhow!("Error parsing code: {}", e.to_string()))?; + + let params = ast.into_iter().find_map(|x| match x { + Stmt::FunctionDef(StmtFunctionDef { name, args, .. }) if &name == &main_name => { + Some(*args) + } + _ => None, + }); + + (params, None) + } + }; + + // Check if main function was found + if params.is_none() { return Ok(MainArgSignature { star_args: false, star_kwargs: false, @@ -241,13 +304,6 @@ pub fn parse_python_signature( has_preprocessor: Some(has_preprocessor), }); } - let ast = Suite::parse(&filtered_code, "main.py") - .map_err(|e| anyhow::anyhow!("Error parsing code: {}", e.to_string()))?; - - let params = ast.into_iter().find_map(|x| match x { - Stmt::FunctionDef(StmtFunctionDef { name, args, .. }) if &name == &main_name => Some(*args), - _ => None, - }); if !skip_params && params.is_some() { let params = params.unwrap(); @@ -259,6 +315,7 @@ pub fn parse_python_signature( // This ensures zero overhead for scripts without enums/docstrings let empty_enums = HashMap::new(); + let module_ref = module.as_ref().map(|m| m.as_slice()); let args_first_pass: Vec<_> = params .args .iter() @@ -269,7 +326,9 @@ pub fn parse_python_signature( .as_arg() .annotation .as_ref() - .map_or((Typ::Unknown, false), |e| parse_expr(e, &empty_enums)); + .map_or((Typ::Unknown, false), |e| { + parse_expr(e, &empty_enums, module_ref) + }); (i, arg_name, typ, has_default) }) .collect(); @@ -282,10 +341,7 @@ pub fn parse_python_signature( let metadata = if has_potential_enums || code.contains("Args:") { extract_code_metadata(code, &main_name) } else { - CodeMetadata { - enums: HashMap::new(), - descriptions: HashMap::new(), - } + CodeMetadata { enums: HashMap::new(), descriptions: HashMap::new() } }; // Build final args, re-parsing Resource types as enums if metadata was extracted @@ -297,7 +353,8 @@ pub fn parse_python_signature( .map(|(i, arg_name, mut typ, mut has_default)| { if matches!(typ, Typ::Resource(_)) && !metadata.enums.is_empty() { if let Some(annotation) = params.args[i].as_arg().annotation.as_ref() { - (typ, has_default) = parse_expr(annotation, &metadata.enums); + (typ, has_default) = + parse_expr(annotation, &metadata.enums, module_ref); } } @@ -357,15 +414,19 @@ pub fn parse_python_signature( } } -fn parse_expr(e: &Box, enums: &HashMap) -> (Typ, bool) { +fn parse_expr( + e: &Box, + enums: &HashMap, + module: Option<&[Stmt]>, +) -> (Typ, bool) { match e.as_ref() { - Expr::Name(ExprName { id, .. }) => (parse_typ(id.as_ref(), enums), false), + Expr::Name(ExprName { id, .. }) => (parse_typ(id.as_ref(), enums, module), false), Expr::Attribute(x) => { if x.value .as_name_expr() .is_some_and(|x| x.id.as_str() == "wmill") { - (parse_typ(x.attr.as_str(), enums), false) + (parse_typ(x.attr.as_str(), enums, module), false) } else { (Typ::Unknown, false) } @@ -375,7 +436,7 @@ fn parse_expr(e: &Box, enums: &HashMap) -> (Typ, bool) { x.right.as_ref(), Expr::Constant(ExprConstant { value: Constant::None, .. }) ) { - (parse_expr(&x.left, enums).0, true) + (parse_expr(&x.left, enums, module).0, true) } else { (Typ::Unknown, false) } @@ -404,8 +465,11 @@ fn parse_expr(e: &Box, enums: &HashMap) -> (Typ, bool) { }; (Typ::Str(values), false) } - "List" | "list" => (Typ::List(Box::new(parse_expr(&x.slice, enums).0)), false), - "Optional" => (parse_expr(&x.slice, enums).0, true), + "List" | "list" => ( + Typ::List(Box::new(parse_expr(&x.slice, enums, module).0)), + false, + ), + "Optional" => (parse_expr(&x.slice, enums, module).0, true), _ => (Typ::Unknown, false), }, _ => (Typ::Unknown, false), @@ -414,7 +478,7 @@ fn parse_expr(e: &Box, enums: &HashMap) -> (Typ, bool) { } } -fn parse_typ(id: &str, enums: &HashMap) -> Typ { +fn parse_typ(id: &str, enums: &HashMap, module: Option<&[Stmt]>) -> Typ { if let Some(enum_info) = enums.get(id) { return Typ::Str(Some(enum_info.values.clone())); } @@ -436,7 +500,17 @@ fn parse_typ(id: &str, enums: &HashMap) -> 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) = 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)) + } } } @@ -470,7 +544,10 @@ fn to_value(et: &Expr, enums: &HashMap) -> Option { - let v = elts.into_iter().map(|x| to_value(&x, enums)).collect::>(); + let v = elts + .into_iter() + .map(|x| to_value(&x, enums)) + .collect::>(); Some(json!(v)) } Expr::Attribute(ExprAttribute { value, attr, .. }) => { @@ -987,10 +1064,17 @@ def main(color: Color = Color.RED): assert_eq!(result.args[0].name, "color"); assert_eq!( result.args[0].typ, - Typ::Str(Some(vec!["red".to_string(), "green".to_string(), "blue".to_string()])) + Typ::Str(Some(vec![ + "red".to_string(), + "green".to_string(), + "blue".to_string() + ])) ); assert_eq!(result.args[0].default, Some(json!("red"))); - assert_eq!(result.args[0].otyp, Some("Color selection from Color enum".to_string())); + assert_eq!( + result.args[0].otyp, + Some("Color selection from Color enum".to_string()) + ); Ok(()) } } diff --git a/backend/parsers/windmill-parser-py/src/pydantic_parser.rs b/backend/parsers/windmill-parser-py/src/pydantic_parser.rs new file mode 100644 index 0000000000..cff5598810 --- /dev/null +++ b/backend/parsers/windmill-parser-py/src/pydantic_parser.rs @@ -0,0 +1,363 @@ +/* + * Author: Ruben Fiszel + * Copyright: Windmill Labs, Inc 2022 + * This file and its contents are licensed under the AGPLv3 License. + * Please see the included NOTICE for copyright information and + * LICENSE-AGPL for a copy of the license. + */ + +//! Pydantic BaseModel and Python dataclass detection and parsing. +//! +//! This module provides functionality to detect and parse Pydantic models and Python +//! dataclasses from Python AST, enabling automatic UI generation for complex data structures. + +use rustpython_parser::ast::{ + Constant, Expr, ExprAttribute, ExprCall, ExprConstant, ExprName, ExprTuple, Stmt, StmtAnnAssign, +}; +use std::collections::HashSet; +use windmill_parser::{ObjectProperty, ObjectType, Typ}; + +// ================================================================== +// Constants +// ================================================================== + +/// Maximum number of fields allowed in a Pydantic model or dataclass. +/// This prevents malicious code from defining models with thousands of fields. +const MAX_MODEL_FIELDS: usize = 200; + +/// Maximum recursion depth for nested types. +const MAX_RECURSION_DEPTH: u8 = 10; + +// ================================================================== +// Pydantic/Dataclass Detection +// ================================================================== + +/// Detects if a class name refers to a Pydantic model or dataclass. +/// Returns ObjectType with parsed fields if detected, None otherwise. +/// +/// # Arguments +/// * `class_name` - The name of the class to look up +/// * `module` - The AST statements to search in +pub fn detect_model_type(class_name: &str, module: &[Stmt]) -> Option { + let mut visited = HashSet::new(); + detect_model_type_impl(class_name, module, &mut visited) +} + +/// Internal implementation with visited tracking +fn detect_model_type_impl( + class_name: &str, + module: &[Stmt], + visited: &mut HashSet, +) -> Option { + // Cycle detection: if we're already parsing this class, return a placeholder + if visited.contains(class_name) { + return Some(ObjectType { + name: Some(class_name.to_string()), + props: None, // Placeholder for self-referential types + }); + } + + // Find class definition in module + for stmt in module { + if let Stmt::ClassDef(class_def) = stmt { + if class_def.name.as_str() == class_name { + // Mark as being visited + visited.insert(class_name.to_string()); + + let result = if is_pydantic_base(&class_def.bases) { + // Pydantic BaseModel + parse_model_fields(&class_def.body, class_def.name.as_str(), module, visited) + } else if has_dataclass_decorator(&class_def.decorator_list) + || has_pydantic_dataclass_decorator(&class_def.decorator_list) + { + // Standard dataclass or Pydantic dataclass + parse_model_fields(&class_def.body, class_def.name.as_str(), module, visited) + } else { + // Found class but it's neither Pydantic nor dataclass + None + }; + + // Remove from visited set after processing + visited.remove(class_name); + + return result; + } + } + } + + // Class not found in module + None +} + +/// Checks if a class inherits from BaseModel or pydantic.BaseModel +fn is_pydantic_base(bases: &[Expr]) -> bool { + for base in bases { + match base { + // Match: class User(BaseModel) + Expr::Name(ExprName { id, .. }) if id.as_str() == "BaseModel" => { + return true; + } + + // Match: class User(pydantic.BaseModel) + Expr::Attribute(ExprAttribute { attr, value, .. }) if attr.as_str() == "BaseModel" => { + if let Expr::Name(ExprName { id, .. }) = value.as_ref() { + if id.as_str() == "pydantic" { + return true; + } + } + } + + _ => {} + } + } + + false +} + +/// Checks if a class has @dataclass decorator (standard library) +fn has_dataclass_decorator(decorators: &[Expr]) -> bool { + for decorator in decorators { + match decorator { + // Match: @dataclass + Expr::Name(ExprName { id, .. }) if id.as_str() == "dataclass" => { + return true; + } + + // Match: @dataclasses.dataclass + Expr::Attribute(ExprAttribute { attr, value, .. }) if attr.as_str() == "dataclass" => { + if let Expr::Name(ExprName { id, .. }) = value.as_ref() { + if id.as_str() == "dataclasses" { + return true; + } + } + } + + // Match: @dataclass() or @dataclass(frozen=True) + Expr::Call(ExprCall { func, .. }) => { + if let Expr::Name(ExprName { id, .. }) = func.as_ref() { + if id.as_str() == "dataclass" { + return true; + } + } + + // Also check for @dataclasses.dataclass(...) + if let Expr::Attribute(ExprAttribute { attr, value, .. }) = func.as_ref() { + if attr.as_str() == "dataclass" { + if let Expr::Name(ExprName { id, .. }) = value.as_ref() { + if id.as_str() == "dataclasses" { + return true; + } + } + } + } + } + + _ => {} + } + } + + false +} + +/// Checks if a class has @pydantic.dataclasses.dataclass decorator (Pydantic v2) +fn has_pydantic_dataclass_decorator(decorators: &[Expr]) -> bool { + for decorator in decorators { + match decorator { + // Match: @pydantic.dataclasses.dataclass + Expr::Attribute(ExprAttribute { attr, value, .. }) if attr.as_str() == "dataclass" => { + if let Expr::Attribute(ExprAttribute { + attr: inner_attr, value: inner_value, .. + }) = value.as_ref() + { + if inner_attr.as_str() == "dataclasses" { + if let Expr::Name(ExprName { id, .. }) = inner_value.as_ref() { + if id.as_str() == "pydantic" { + return true; + } + } + } + } + } + + // Match: @pydantic.dataclasses.dataclass(...) + Expr::Call(ExprCall { func, .. }) => { + if let Expr::Attribute(ExprAttribute { attr, value, .. }) = func.as_ref() { + if attr.as_str() == "dataclass" { + if let Expr::Attribute(ExprAttribute { + attr: inner_attr, + value: inner_value, + .. + }) = value.as_ref() + { + if inner_attr.as_str() == "dataclasses" { + if let Expr::Name(ExprName { id, .. }) = inner_value.as_ref() { + if id.as_str() == "pydantic" { + return true; + } + } + } + } + } + } + } + + _ => {} + } + } + + false +} + +// ================================================================== +// Field Parsing (Unified for Pydantic and Dataclass) +// ================================================================== + +/// Parses model fields from class body (works for both Pydantic and dataclass) +fn parse_model_fields( + body: &[Stmt], + class_name: &str, + module: &[Stmt], + visited: &mut HashSet, +) -> Option { + let mut properties = Vec::new(); + + for stmt in body { + // Extract annotated assignments: field_name: field_type + if let Stmt::AnnAssign(ann_assign) = stmt { + if let Some(prop) = parse_annotated_field(ann_assign, module, visited) { + if properties.len() >= MAX_MODEL_FIELDS { + eprintln!( + "Model {model} exceeds maximum field count {limit}, truncating", + model = class_name, + limit = MAX_MODEL_FIELDS.to_string(), + ); + break; + } + properties.push(prop); + } + } + } + + if properties.is_empty() { + // Empty model - return object with no properties + return Some(ObjectType { name: Some(class_name.to_string()), props: None }); + } + + Some(ObjectType { name: Some(class_name.to_string()), props: Some(properties) }) +} + +/// Parses a single annotated field assignment +fn parse_annotated_field( + ann_assign: &StmtAnnAssign, + module: &[Stmt], + visited: &mut HashSet, +) -> Option { + if let Expr::Name(ExprName { id: field_name, .. }) = ann_assign.target.as_ref() { + let field_type = extract_field_type(&ann_assign.annotation, 0, module, visited); + Some(ObjectProperty { key: field_name.to_string(), typ: Box::new(field_type) }) + } else { + None + } +} + +// ================================================================== +// Type Extraction +// ================================================================== + +/// Extracts Windmill Typ from Python type annotation (RECURSIVE). +/// +/// # Arguments +/// * `annotation` - The Python AST expression representing the type annotation +/// * `depth` - Current recursion depth (prevents infinite recursion) +/// * `module` - The AST statements for nested model lookup +/// * `visited` - Set of class names currently being parsed (for cycle detection) +fn extract_field_type( + annotation: &Expr, + depth: u8, + module: &[Stmt], + visited: &mut HashSet, +) -> Typ { + // Prevent infinite recursion + if depth >= MAX_RECURSION_DEPTH { + eprintln!( + "Type annotation recursion limit {limit} reached, returning Unknown type", + limit = MAX_RECURSION_DEPTH, + ); + return Typ::Unknown; + } + + match annotation { + // Simple types: str, int, bool, float, bytes, Any + Expr::Name(ExprName { id, .. }) => match id.as_str() { + "str" => Typ::Str(None), + "int" => Typ::Int, + "float" => Typ::Float, + "bool" => Typ::Bool, + "bytes" => Typ::Bytes, + "datetime" => Typ::Datetime, + "Any" => Typ::Unknown, // typing.Any maps to Unknown + // Custom class - check if it's a model + custom_type => { + if let Some(object_type) = detect_model_type_impl(custom_type, module, visited) { + Typ::Object(object_type) + } else { + // Unknown type - return Unknown instead of Resource + Typ::Unknown + } + } + }, + + // Generic types: List[T], Optional[T], Dict[K, V], Annotated[T, ...] + Expr::Subscript(subscript) => { + if let Expr::Name(ExprName { id, .. }) = subscript.value.as_ref() { + match id.as_str() { + // List[T] + "List" | "list" => { + let inner_type = + extract_field_type(&subscript.slice, depth + 1, module, visited); + Typ::List(Box::new(inner_type)) + } + + // Optional[T] - unwrap to T + "Optional" => extract_field_type(&subscript.slice, depth + 1, module, visited), + + // Dict[K, V] - return generic Object + "Dict" | "dict" => Typ::Object(ObjectType::new(None, Some(vec![]))), + + // Annotated[T, ...] - extract the first type argument (Pydantic v2) + "Annotated" => match subscript.slice.as_ref() { + Expr::Tuple(ExprTuple { elts, .. }) if !elts.is_empty() => { + extract_field_type(&elts[0], depth + 1, module, visited) + } + _ => Typ::Unknown, + }, + + // Set[T], Tuple[T], etc. - not supported + _ => Typ::Unknown, + } + } else { + Typ::Unknown + } + } + + // Union types: str | int (Python 3.10+) or Union[str, int] + // Not fully supported - return Unknown with warning + Expr::BinOp(_) => { + eprintln!("Union types (e.g., str | int) are not yet supported, treating as Unknown"); + Typ::Unknown + } + + // String annotations: "ForwardRef" (forward references) + // Not fully supported - return Unknown with warning + Expr::Constant(ExprConstant { value: Constant::Str(s), .. }) => { + eprintln!( + "Forward references like \"{forward_ref}\" are not yet supported, treating as Unknown", + forward_ref = s, + ); + Typ::Unknown + } + + // All other annotations + _ => Typ::Unknown, + } +} diff --git a/backend/parsers/windmill-parser-py/tests/pydantic_tests.rs b/backend/parsers/windmill-parser-py/tests/pydantic_tests.rs new file mode 100644 index 0000000000..89f66684d7 --- /dev/null +++ b/backend/parsers/windmill-parser-py/tests/pydantic_tests.rs @@ -0,0 +1,699 @@ +/* + * Author: Ruben Fiszel + * Copyright: Windmill Labs, Inc 2022 + * This file and its contents are licensed under the AGPLv3 License. + * Please see the included NOTICE for copyright information and + * LICENSE-AGPL for a copy of the license. + */ + +//! Integration tests for Pydantic BaseModel and Python dataclass support. + +use windmill_parser::Typ; +use windmill_parser_py::parse_python_signature; + +#[test] +fn test_pydantic_basic_model() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel + +class User(BaseModel): + name: str + age: int + email: str + +def main(user: User): + return f'Hello {user.name}' +"; + let result = parse_python_signature(code, None, false)?; + + // Check that user parameter is detected as Object type + assert_eq!(result.args.len(), 1); + assert_eq!(result.args[0].name, "user"); + + // Verify it's an Object type with correct model name + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("User".to_string())); + assert!(obj.props.is_some()); + + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 3); + + // Verify field names and types + assert_eq!(props[0].key, "name"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + assert_eq!(props[1].key, "age"); + assert_eq!(*props[1].typ, Typ::Int); + + assert_eq!(props[2].key, "email"); + assert_eq!(*props[2].typ, Typ::Str(None)); + } + _ => panic!("Expected Typ::Object for Pydantic model"), + } + + Ok(()) +} + +#[test] +fn test_python_dataclass() -> anyhow::Result<()> { + let code = " +from dataclasses import dataclass + +@dataclass +class Config: + host: str + port: int + debug: bool + +def main(config: Config): + return config.host +"; + let result = parse_python_signature(code, None, false)?; + + // Check that config parameter is detected as Object type + assert_eq!(result.args.len(), 1); + assert_eq!(result.args[0].name, "config"); + + // Verify it's an Object type with correct class name + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("Config".to_string())); + assert!(obj.props.is_some()); + + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 3); + + // Verify field names and types + assert_eq!(props[0].key, "host"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + assert_eq!(props[1].key, "port"); + assert_eq!(*props[1].typ, Typ::Int); + + assert_eq!(props[2].key, "debug"); + assert_eq!(*props[2].typ, Typ::Bool); + } + _ => panic!("Expected Typ::Object for dataclass"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_nested_model() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel + +class Address(BaseModel): + street: str + city: str + +class Person(BaseModel): + name: str + address: Address + +def main(person: Person): + return person.name +"; + let result = parse_python_signature(code, None, false)?; + + // Check that person parameter is detected as Object type + assert_eq!(result.args.len(), 1); + assert_eq!(result.args[0].name, "person"); + + // Verify it's an Object type with nested model + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("Person".to_string())); + assert!(obj.props.is_some()); + + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + + // Verify name field + assert_eq!(props[0].key, "name"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + // Verify address field is a nested Object + assert_eq!(props[1].key, "address"); + match props[1].typ.as_ref() { + Typ::Object(nested_obj) => { + assert_eq!(nested_obj.name, Some("Address".to_string())); + assert!(nested_obj.props.is_some()); + + let nested_props = nested_obj.props.as_ref().unwrap(); + assert_eq!(nested_props.len(), 2); + assert_eq!(nested_props[0].key, "street"); + assert_eq!(nested_props[1].key, "city"); + } + _ => panic!("Expected nested Typ::Object for Address"), + } + } + _ => panic!("Expected Typ::Object for Person model"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_empty_model() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel + +class EmptyModel(BaseModel): + pass + +def main(model: EmptyModel): + return 'ok' +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("EmptyModel".to_string())); + assert!(obj.props.is_none()); + } + _ => panic!("Expected Typ::Object for empty model"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_list_field() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel +from typing import List + +class TodoList(BaseModel): + items: List[str] + count: int + +def main(todos: TodoList): + return todos.count +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("TodoList".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + + // Verify List[str] type + assert_eq!(props[0].key, "items"); + match props[0].typ.as_ref() { + Typ::List(inner) => { + assert_eq!(**inner, Typ::Str(None)); + } + _ => panic!("Expected Typ::List for items field"), + } + + assert_eq!(props[1].key, "count"); + assert_eq!(*props[1].typ, Typ::Int); + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_optional_field() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel +from typing import Optional + +class User(BaseModel): + name: str + nickname: Optional[str] + +def main(user: User): + return user.name +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + + assert_eq!(props[0].key, "name"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + // Optional[str] should unwrap to str + assert_eq!(props[1].key, "nickname"); + assert_eq!(*props[1].typ, Typ::Str(None)); + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_dataclass_with_decorator_args() -> anyhow::Result<()> { + let code = " +from dataclasses import dataclass + +@dataclass(frozen=True) +class ImmutableConfig: + setting: str + value: int + +def main(config: ImmutableConfig): + return config.setting +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("ImmutableConfig".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + } + _ => panic!("Expected Typ::Object for dataclass"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_dict_field() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel +from typing import Dict + +class Config(BaseModel): + settings: Dict[str, str] + name: str + +def main(config: Config): + return config.name +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + + // Dict should return generic Object + assert_eq!(props[0].key, "settings"); + match props[0].typ.as_ref() { + Typ::Object(_) => {} // Generic object for Dict + _ => panic!("Expected Typ::Object for Dict field"), + } + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_non_model_class_treated_as_resource() -> anyhow::Result<()> { + let code = " +class RegularClass: + def __init__(self, value): + self.value = value + +def main(obj: RegularClass): + return 'ok' +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + // Regular classes (non-Pydantic/dataclass) should be treated as Resource + assert_eq!( + result.args[0].typ, + Typ::Resource("RegularClass".to_string()) + ); + + Ok(()) +} + +#[test] +fn test_invalid_syntax_fallback() -> anyhow::Result<()> { + // Code with syntax errors - should still not crash + let code = " +from pydantic import BaseModel + +class User(BaseModel: # Missing closing paren + name: str + +def main(user: User): + return 'ok' +"; + // Should not panic, even with invalid syntax + let result = parse_python_signature(code, None, false); + + // Either succeeds with Unknown types or fails gracefully + assert!(result.is_ok() || result.is_err()); + + Ok(()) +} + +#[test] +fn test_datetime_type() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel +from datetime import datetime + +class Event(BaseModel): + name: str + created_at: datetime + +def main(event: Event): + return event.name +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + + assert_eq!(props[0].key, "name"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + assert_eq!(props[1].key, "created_at"); + assert_eq!(*props[1].typ, Typ::Datetime); + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_multiple_pydantic_models() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel + +class User(BaseModel): + name: str + +class Post(BaseModel): + title: str + author: User + +def main(post: Post): + return post.title +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("Post".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + + // Nested User model + assert_eq!(props[1].key, "author"); + match props[1].typ.as_ref() { + Typ::Object(nested) => { + assert_eq!(nested.name, Some("User".to_string())); + } + _ => panic!("Expected nested Typ::Object for User"), + } + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_self_referential_model() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel +from typing import List, Optional + +class TreeNode(BaseModel): + value: str + children: List[TreeNode] + parent: Optional[TreeNode] + +def main(root: TreeNode): + return root.value +"; + let result = parse_python_signature(code, None, false)?; + + // Should not panic and handle the cycle gracefully + assert_eq!(result.args.len(), 1); + assert_eq!(result.args[0].name, "root"); + + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("TreeNode".to_string())); + assert!(obj.props.is_some()); + + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 3); + + // value: str + assert_eq!(props[0].key, "value"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + // children: List[TreeNode] - self-reference should return placeholder + assert_eq!(props[1].key, "children"); + match props[1].typ.as_ref() { + Typ::List(inner) => match inner.as_ref() { + Typ::Object(nested) => { + assert_eq!(nested.name, Some("TreeNode".to_string())); + // Placeholder has no props (to break the cycle) + assert!(nested.props.is_none()); + } + _ => panic!("Expected nested Typ::Object for TreeNode"), + }, + _ => panic!("Expected Typ::List for children"), + } + + // parent: Optional[TreeNode] - self-reference should return placeholder + assert_eq!(props[2].key, "parent"); + match props[2].typ.as_ref() { + Typ::Object(nested) => { + assert_eq!(nested.name, Some("TreeNode".to_string())); + assert!(nested.props.is_none()); + } + _ => panic!("Expected Typ::Object for parent"), + } + } + _ => panic!("Expected Typ::Object for TreeNode"), + } + + Ok(()) +} + +#[test] +fn test_any_type() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel +from typing import Any + +class FlexibleModel(BaseModel): + name: str + data: Any + metadata: Any + +def main(model: FlexibleModel): + return model.name +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("FlexibleModel".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 3); + + assert_eq!(props[0].key, "name"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + // Any should map to Unknown + assert_eq!(props[1].key, "data"); + assert_eq!(*props[1].typ, Typ::Unknown); + + assert_eq!(props[2].key, "metadata"); + assert_eq!(*props[2].typ, Typ::Unknown); + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_annotated_type() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel, Field +from typing import Annotated + +class User(BaseModel): + name: Annotated[str, Field(min_length=1)] + age: Annotated[int, Field(ge=0)] + email: Annotated[str, Field(pattern=r'^[a-z]+@[a-z]+\\.[a-z]+$')] + +def main(user: User): + return user.name +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("User".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 3); + + // Annotated[str, ...] should unwrap to str + assert_eq!(props[0].key, "name"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + // Annotated[int, ...] should unwrap to int + assert_eq!(props[1].key, "age"); + assert_eq!(*props[1].typ, Typ::Int); + + // Annotated[str, ...] should unwrap to str + assert_eq!(props[2].key, "email"); + assert_eq!(*props[2].typ, Typ::Str(None)); + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_dataclass() -> anyhow::Result<()> { + let code = " +import pydantic.dataclasses + +@pydantic.dataclasses.dataclass +class PydanticConfig: + host: str + port: int + debug: bool + +def main(config: PydanticConfig): + return config.host +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + assert_eq!(result.args[0].name, "config"); + + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("PydanticConfig".to_string())); + assert!(obj.props.is_some()); + + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 3); + + assert_eq!(props[0].key, "host"); + assert_eq!(*props[0].typ, Typ::Str(None)); + + assert_eq!(props[1].key, "port"); + assert_eq!(*props[1].typ, Typ::Int); + + assert_eq!(props[2].key, "debug"); + assert_eq!(*props[2].typ, Typ::Bool); + } + _ => panic!("Expected Typ::Object for pydantic dataclass"), + } + + Ok(()) +} + +#[test] +fn test_pydantic_dataclass_with_args() -> anyhow::Result<()> { + let code = " +import pydantic.dataclasses + +@pydantic.dataclasses.dataclass(frozen=True) +class ImmutablePydanticConfig: + name: str + value: int + +def main(config: ImmutablePydanticConfig): + return config.name +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("ImmutablePydanticConfig".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 2); + } + _ => panic!("Expected Typ::Object for pydantic dataclass"), + } + + Ok(()) +} + +#[test] +fn test_unknown_type_in_pydantic_field() -> anyhow::Result<()> { + let code = " +from pydantic import BaseModel + +class SomeOtherClass: + pass + +class Model(BaseModel): + field: SomeOtherClass + +def main(m: Model): + return 'ok' +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 1); + match &result.args[0].typ { + Typ::Object(obj) => { + assert_eq!(obj.name, Some("Model".to_string())); + let props = obj.props.as_ref().unwrap(); + assert_eq!(props.len(), 1); + + // SomeOtherClass inside Model becomes Unknown (not Resource) + assert_eq!(props[0].key, "field"); + assert_eq!(*props[0].typ, Typ::Unknown); + } + _ => panic!("Expected Typ::Object"), + } + + Ok(()) +} + +#[test] +fn test_simple_script_without_models() -> anyhow::Result<()> { + // This test verifies the optimization: simple scripts without Pydantic/dataclass + // should not trigger the expensive full AST parse + let code = " +def main(name: str, age: int, active: bool = True): + return f'Hello {name}, you are {age} years old' +"; + let result = parse_python_signature(code, None, false)?; + + assert_eq!(result.args.len(), 3); + assert_eq!(result.args[0].name, "name"); + assert_eq!(result.args[0].typ, Typ::Str(None)); + assert_eq!(result.args[1].name, "age"); + assert_eq!(result.args[1].typ, Typ::Int); + assert_eq!(result.args[2].name, "active"); + assert_eq!(result.args[2].typ, Typ::Bool); + + Ok(()) +} diff --git a/cli/wasm/py/windmill_parser_wasm_bg.wasm b/cli/wasm/py/windmill_parser_wasm_bg.wasm index 3d33d436c5..eb896993f0 100644 Binary files a/cli/wasm/py/windmill_parser_wasm_bg.wasm and b/cli/wasm/py/windmill_parser_wasm_bg.wasm differ diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 3b56a7fc1a..3c9b0969d3 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -79,7 +79,7 @@ "windmill-parser-wasm-java": "1.510.1", "windmill-parser-wasm-nu": "1.510.1", "windmill-parser-wasm-php": "1.574.1", - "windmill-parser-wasm-py": "1.595.0", + "windmill-parser-wasm-py": "1.601.1", "windmill-parser-wasm-regex": "1.593.0", "windmill-parser-wasm-ruby": "1.526.1", "windmill-parser-wasm-rust": "1.558.1", @@ -15858,9 +15858,9 @@ "integrity": "sha512-COyid6B1RYs+bpzUCInsA4HY/WZkpDLfkQ90+AqU/TVTpzYSbAC2JCbIwy0cRElBvlhI4bQ+9Wg6hSQKMpEkpA==" }, "node_modules/windmill-parser-wasm-py": { - "version": "1.595.0", - "resolved": "https://registry.npmjs.org/windmill-parser-wasm-py/-/windmill-parser-wasm-py-1.595.0.tgz", - "integrity": "sha512-jkZiEl43J7jyaQDSckMEdSr0hgG3pFVt7pT8MbfNFHV2ywxnkdQsk8fwfdIo3rHHrErPPd4OVRxq0kulCTClXg==" + "version": "1.601.1", + "resolved": "https://registry.npmjs.org/windmill-parser-wasm-py/-/windmill-parser-wasm-py-1.601.1.tgz", + "integrity": "sha512-xcNZE/8B29yfl6UuQDPSXMD+83/W2Hzt2uhn+WrNvy0+qzk6nLh/vJGrf2srLBngYX1TxhUI5Jgseg0PK9yvNw==" }, "node_modules/windmill-parser-wasm-regex": { "version": "1.593.0", diff --git a/frontend/package.json b/frontend/package.json index 593a348b8e..b5dd310736 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -151,7 +151,7 @@ "windmill-parser-wasm-java": "1.510.1", "windmill-parser-wasm-nu": "1.510.1", "windmill-parser-wasm-php": "1.574.1", - "windmill-parser-wasm-py": "1.595.0", + "windmill-parser-wasm-py": "1.601.1", "windmill-parser-wasm-regex": "1.593.0", "windmill-parser-wasm-ruby": "1.526.1", "windmill-parser-wasm-rust": "1.558.1",