From d09952572cfa004a36a3dcd9f81669cc0bc7bf63 Mon Sep 17 00:00:00 2001 From: Avigail Royzenberg <58378974+AvigailRoyzenberg@users.noreply.github.com> Date: Fri, 19 Dec 2025 15:41:34 -0500 Subject: [PATCH] fix: Python Enum types generate proper dropdown schemas with descriptions (#7400) * Fix Python Enum and Literal schema generation with docstring descriptions - Extract Enum class definitions and their string values - Parse docstring Args: sections for parameter descriptions - Map Enum type annotations to string enums with proper values - Handle Enum.VALUE default values correctly - Store descriptions in Arg.otyp field - Add test case for enum with docstring parsing * perf: optimize enum parser and fix default value handling - Combine enum extraction and docstring parsing into single AST pass (2x performance improvement) - Add support for IntEnum, StrEnum, Flag, IntFlag types - Fix default values to use actual enum values (e.g., 'red') instead of member names (e.g., 'RED') - Improve docstring parsing robustness with proper indentation tracking - Clean up code structure with EnumInfo type for better maintainability All tests pass. This addresses code review feedback for performance and correctness. * perf: implement true lazy evaluation for enum parsing - Only parse metadata when unknown types encountered - Two-pass approach: parse types first, extract only if needed - Zero overhead for scripts without enums - Keyword checks + prepass filtering when extraction needed --- backend/parsers/windmill-parser-py/src/lib.rs | 294 ++++++++++++++++-- 1 file changed, 268 insertions(+), 26 deletions(-) diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index 002ea51cbc..4c45b743e5 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/backend/parsers/windmill-parser-py/src/lib.rs @@ -15,7 +15,7 @@ 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, ExprAttribute, ExprConstant, ExprDict, ExprList, ExprName, Stmt, StmtAssign, StmtClassDef, StmtFunctionDef, Suite, }, Parse, }; @@ -60,6 +60,166 @@ fn filter_non_main(code: &str, main_name: &str) -> String { return filtered_code; } +/// Data extracted from parsing the Python code +struct CodeMetadata { + enums: HashMap, + descriptions: HashMap, +} + +/// Information about an Enum class +struct EnumInfo { + values: Vec, + members: HashMap, +} + +fn has_enum_keyword(code: &str) -> bool { + code.contains("Enum") +} + +/// Extract only class and function definitions from code (prepass filtering) +fn filter_relevant_statements(code: &str) -> String { + let mut result = Vec::new(); + let mut lines = code.lines().peekable(); + + while let Some(line) = lines.next() { + let trimmed = line.trim_start(); + + if trimmed.starts_with("class ") || trimmed.starts_with("def ") { + result.push(line); + let base_indent = line.len() - trimmed.len(); + + while let Some(&next_line) = lines.peek() { + let next_trimmed = next_line.trim_start(); + let next_indent = next_line.len() - next_trimmed.len(); + + if next_trimmed.is_empty() || next_indent > base_indent { + result.push(lines.next().unwrap()); + } else { + break; + } + } + } + } + + result.join("\n") +} + +/// Extract Enum definitions and docstring descriptions lazily. +/// Only parses AST if relevant keywords are present. +fn extract_code_metadata(code: &str, main_name: &str) -> CodeMetadata { + let mut enums = HashMap::new(); + let mut descriptions = HashMap::new(); + + let has_enum = has_enum_keyword(code); + let has_docstring = code.contains("Args:"); + + if !has_enum && !has_docstring { + return CodeMetadata { enums, descriptions }; + } + + let filtered_code = filter_relevant_statements(code); + + let ast = match Suite::parse(&filtered_code, "main.py") { + Ok(ast) => ast, + Err(_) => return CodeMetadata { enums, descriptions }, + }; + + for stmt in ast { + match stmt { + Stmt::ClassDef(StmtClassDef { name, body, bases, .. }) if has_enum => { + let is_enum = bases.iter().any(|base| { + matches!(base, Expr::Name(ExprName { id, .. }) + if id == "Enum" || id == "IntEnum" || id == "StrEnum" + || id == "Flag" || id == "IntFlag") + }); + + if is_enum { + let mut values = Vec::new(); + let mut members = HashMap::new(); + + for item in body { + if let Stmt::Assign(StmtAssign { targets, value, .. }) = item { + 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() { + values.push(val.to_string()); + members.insert(target_name.to_string(), val.to_string()); + } + } + } + } + } + + if !values.is_empty() { + 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() { + descriptions = parse_docstring_args(docstring); + } + } + } + }, + _ => {} + } + } + + CodeMetadata { enums, descriptions } +} + +/// Parse docstring Args: section (format: "param_name (type): Description") +fn parse_docstring_args(docstring: &str) -> HashMap { + let mut descriptions = HashMap::new(); + let mut in_args_section = false; + let mut base_indent: Option = None; + + for line in docstring.lines() { + let trimmed = line.trim(); + + if trimmed == "Args:" { + in_args_section = true; + base_indent = None; + continue; + } + + if in_args_section { + if trimmed.is_empty() { + continue; + } + + let indent = line.len() - line.trim_start().len(); + + if base_indent.is_none() && !trimmed.is_empty() { + base_indent = Some(indent); + } + + if let Some(base) = base_indent { + if indent < base && trimmed.ends_with(':') { + break; + } + } + + if let Some(colon_pos) = trimmed.find(':') { + let before_colon = &trimmed[..colon_pos]; + let description = trimmed[colon_pos + 1..].trim(); + + if let Some(paren_pos) = before_colon.find('(') { + let param_name = before_colon[..paren_pos].trim(); + descriptions.insert(param_name.to_string(), description.to_string()); + } else { + descriptions.insert(before_colon.trim().to_string(), description.to_string()); + } + } + } + } + + descriptions +} + /// skip_params is a micro optimization for when we just want to find the main /// function without parsing all the params. pub fn parse_python_signature( @@ -91,27 +251,61 @@ pub fn parse_python_signature( if !skip_params && params.is_some() { let params = params.unwrap(); - //println!("{:?}", params); let def_arg_start = params.args.len() - params.defaults().count(); + + // Two-pass approach for lazy metadata extraction: + // Pass 1: Parse types without enum info to determine if metadata is needed + // Pass 2: Re-parse unknown types with metadata only if necessary + // This ensures zero overhead for scripts without enums/docstrings + + let empty_enums = HashMap::new(); + let args_first_pass: Vec<_> = params + .args + .iter() + .enumerate() + .map(|(i, x)| { + let arg_name = x.as_arg().arg.to_string(); + let (typ, has_default) = x + .as_arg() + .annotation + .as_ref() + .map_or((Typ::Unknown, false), |e| parse_expr(e, &empty_enums)); + (i, arg_name, typ, has_default) + }) + .collect(); + + // Determine if we need to extract metadata from the code + let has_potential_enums = args_first_pass + .iter() + .any(|(_, _, typ, _)| matches!(typ, Typ::Resource(_))); + + let metadata = if has_potential_enums || code.contains("Args:") { + extract_code_metadata(code, &main_name) + } else { + CodeMetadata { + enums: HashMap::new(), + descriptions: HashMap::new(), + } + }; + + // Build final args, re-parsing Resource types as enums if metadata was extracted Ok(MainArgSignature { star_args: params.vararg.is_some(), star_kwargs: params.kwarg.is_some(), - args: params - .args - .iter() - .enumerate() - .map(|(i, x)| { - let (mut typ, has_default) = x - .as_arg() - .annotation - .as_ref() - .map_or((Typ::Unknown, false), |e| parse_expr(e)); + args: args_first_pass + .into_iter() + .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); + } + } let default = if i >= def_arg_start { params .defaults() .nth(i - def_arg_start) - .map(to_value) + .map(|expr| to_value(expr, &metadata.enums)) .flatten() } else { None @@ -140,8 +334,8 @@ pub fn parse_python_signature( } Arg { - otyp: None, - name: x.as_arg().arg.to_string(), + otyp: metadata.descriptions.get(&arg_name).map(|d| d.to_string()), + name: arg_name, typ, has_default: has_default || default.is_some(), default, @@ -163,15 +357,15 @@ pub fn parse_python_signature( } } -fn parse_expr(e: &Box) -> (Typ, bool) { +fn parse_expr(e: &Box, enums: &HashMap) -> (Typ, bool) { match e.as_ref() { - Expr::Name(ExprName { id, .. }) => (parse_typ(id.as_ref()), false), + Expr::Name(ExprName { id, .. }) => (parse_typ(id.as_ref(), enums), 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()), false) + (parse_typ(x.attr.as_str(), enums), false) } else { (Typ::Unknown, false) } @@ -181,7 +375,7 @@ fn parse_expr(e: &Box) -> (Typ, bool) { x.right.as_ref(), Expr::Constant(ExprConstant { value: Constant::None, .. }) ) { - (parse_expr(&x.left).0, true) + (parse_expr(&x.left, enums).0, true) } else { (Typ::Unknown, false) } @@ -210,8 +404,8 @@ fn parse_expr(e: &Box) -> (Typ, bool) { }; (Typ::Str(values), false) } - "List" | "list" => (Typ::List(Box::new(parse_expr(&x.slice).0)), false), - "Optional" => (parse_expr(&x.slice).0, true), + "List" | "list" => (Typ::List(Box::new(parse_expr(&x.slice, enums).0)), false), + "Optional" => (parse_expr(&x.slice, enums).0, true), _ => (Typ::Unknown, false), }, _ => (Typ::Unknown, false), @@ -220,7 +414,11 @@ fn parse_expr(e: &Box) -> (Typ, bool) { } } -fn parse_typ(id: &str) -> Typ { +fn parse_typ(id: &str, enums: &HashMap) -> Typ { + if let Some(enum_info) = enums.get(id) { + return Typ::Str(Some(enum_info.values.clone())); + } + match id { "str" => Typ::Str(None), "float" => Typ::Float, @@ -249,7 +447,7 @@ fn map_resource_name(x: &str) -> String { } } -fn to_value(et: &Expr) -> Option { +fn to_value(et: &Expr, enums: &HashMap) -> Option { match et { Expr::Constant(ExprConstant { value, .. }) => Some(constant_to_value(value)), Expr::Dict(ExprDict { keys, values, .. }) => { @@ -259,22 +457,35 @@ fn to_value(et: &Expr) -> Option { .map(|(k, v)| { let key = k .as_ref() - .map(to_value) + .map(|e| to_value(e, enums)) .flatten() .and_then(|x| match x { serde_json::Value::String(s) => Some(s), _ => None, }) .unwrap_or_else(|| "no_key".to_string()); - (key, to_value(&v)) + (key, to_value(&v, enums)) }) .collect::>(); Some(json!(v)) } Expr::List(ExprList { elts, .. }) => { - let v = elts.into_iter().map(|x| to_value(&x)).collect::>(); + let v = elts.into_iter().map(|x| to_value(&x, enums)).collect::>(); Some(json!(v)) } + Expr::Attribute(ExprAttribute { value, attr, .. }) => { + // Handle Enum.MEMBER: returns enum value ("red") not member name ("RED") + if let Expr::Name(ExprName { id: enum_name, .. }) = value.as_ref() { + if let Some(enum_info) = enums.get(enum_name.as_str()) { + if let Some(enum_value) = enum_info.members.get(attr.as_str()) { + return Some(json!(enum_value)); + } + } + Some(json!(attr.as_str())) + } else { + None + } + } Expr::Call { .. } => Some(json!(FUNCTION_CALL)), _ => None, } @@ -751,4 +962,35 @@ def main(a: str, b: Optional[str], c: str | None): return Ok(()) } + + #[test] + fn test_parse_python_sig_enum() -> anyhow::Result<()> { + let code = r#" +from enum import Enum + +class Color(str, Enum): + RED = 'red' + GREEN = 'green' + BLUE = 'blue' + +def main(color: Color = Color.RED): + """ + Test enum parsing + + Args: + color (Color): Color selection from Color enum + """ + return {"color": color} +"#; + let result = parse_python_signature(code, None, false)?; + assert_eq!(result.args.len(), 1); + 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()])) + ); + assert_eq!(result.args[0].default, Some(json!("red"))); + assert_eq!(result.args[0].otyp, Some("Color selection from Color enum".to_string())); + Ok(()) + } }