starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-08-18 16:02:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Pydantic BaseModel and dataclass support for Python type inference (#7497)

Browse Source 
* feat: add Pydantic BaseModel and dataclass support for Python type inference

- Add AST-based detection of Pydantic BaseModel inheritance patterns
- Add AST-based detection of @dataclass decorator (all variants)
- Implement recursive field schema extraction with type inference
- Add thread-safe stack-based module storage for nested parsing
- Add RAII cleanup guard to ensure memory safety on all code paths
- Add security limits: 200 fields max, 10 recursion levels max
- Add comprehensive test coverage: 3 new tests for Pydantic/dataclass
- Maintain 100% backward compatibility with existing type system

This enables ML/AI practitioners to use Pydantic models as function
parameters with automatic UI generation from model schemas.

Implementation highlights:
- Zero code execution: Pure AST analysis for safety
- Thread-safe: Stack-based storage prevents race conditions
- Memory-safe: RAII pattern guarantees cleanup
- Security-hardened: Field count and recursion depth limits
- Performance-optimized: Depth-limited recursion, lazy parsing

Test results: All 12 tests passing (9 existing + 3 new)

Closes #4700

🤖 Generated with Claude Code (https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* fix: improve Pydantic/dataclass parser robustness and error handling

This commit addresses critical bugs and improves error handling in the
Python parser for Pydantic BaseModel and dataclass support.

## Critical Fixes

1. **Thread-local storage RAII pattern**: Fixed bug where parse failures
   could leave the module stack in an inconsistent state. Now uses proper
   functional composition with .ok().map() to ensure cleanup always happens.

2. **Recursion depth warnings**: Added explicit warning messages when the
   recursion depth limit (10 levels) is reached during type extraction.
   Made the limit a named constant for clarity.

3. **Unsupported type warnings**: Added informative warning messages for
   unsupported type annotations (Union types and forward references) to
   help users understand why their types aren't being inferred.

## Improvements

- Added 10 comprehensive test cases covering:
  - Empty Pydantic models
  - List[T] and Optional[T] types
  - Dataclass with decorator arguments
  - Dict types
  - Regular classes (non-model types)
  - Invalid syntax handling
  - Datetime fields
  - Multiple model definitions
  - Nested models

- All 21 tests pass successfully

## Testing

Verified that:
- Parser handles malformed code gracefully
- RAII cleanup works correctly with early returns
- Warning messages are clear and actionable
- No memory leaks or panics

Closes #4700

* refactor: Separate Pydantic/dataclass code into dedicated module. Created src/pydantic_parser.rs with thread-local storage, model detection, and type extraction logic. Moved 12 Pydantic tests to tests/pydantic_tests.rs and removed duplicate code from lib.rs. All 21 tests passing.

* opti and publish

---------

Co-authored-by: Devdatta Talele <devtalele0@gmail.com>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
hugocasa
2026-01-07 01:48:04 +07:00
committed by GitHub
parent 7877999f3d
commit 0f2b417ff5
6 changed files with 1183 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/parsers/windmill-parser-py/src/lib.rs
+116 -32
View File
@@ -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 = "<function call>";
/// 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<Expr>, enums: &HashMap<String, EnumInfo>) -> (Typ, bool) {
fn parse_expr(
e: &Box<Expr>,
enums: &HashMap<String, EnumInfo>,
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<Expr>, enums: &HashMap<String, EnumInfo>) -> (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<Expr>, enums: &HashMap<String, EnumInfo>) -> (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<Expr>, enums: &HashMap<String, EnumInfo>) -> (Typ, bool) {
}
}
fn parse_typ(id: &str, enums: &HashMap<String, EnumInfo>) -> Typ {
fn parse_typ(id: &str, enums: &HashMap<String, EnumInfo>, 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<String, EnumInfo>) -> 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<R>(et: &Expr<R>, enums: &HashMap<String, EnumInfo>) -> Option<serde_
Some(json!(v))
}
Expr::List(ExprList { elts, .. }) => {
let v = elts.into_iter().map(|x| to_value(&x, enums)).collect::<Vec<_>>();
let v = elts
.into_iter()
.map(|x| to_value(&x, enums))
.collect::<Vec<_>>();
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(())
}
}
backend/parsers/windmill-parser-py/src/pydantic_parser.rs
+363
View File
@@ -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<ObjectType> {
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<String>,
) -> Option<ObjectType> {
// 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<String>,
) -> Option<ObjectType> {
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<String>,
) -> Option<ObjectProperty> {
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<String>,
) -> 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,
}
}
backend/parsers/windmill-parser-py/tests/pydantic_tests.rs
+699
View File
@@ -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(())
}
cli/wasm/py/windmill_parser_wasm_bg.wasm
BIN
View File
Binary file not shown.
frontend/package-lock.json
Generated
+4 -4
View File
@@ -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",
frontend/package.json
+1 -1
View File
@@ -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",