diff --git a/mlf-codegen/src/lib.rs b/mlf-codegen/src/lib.rs index c0a0f24..cf65b3d 100644 --- a/mlf-codegen/src/lib.rs +++ b/mlf-codegen/src/lib.rs @@ -72,6 +72,46 @@ fn has_main_annotation(annotations: &[Annotation]) -> bool { annotations.iter().any(|ann| ann.name.name == "main") } +fn get_annotation_string_value(annotations: &[Annotation], name: &str) -> Option { + annotations.iter() + .find(|ann| ann.name.name == name) + .and_then(|ann| { + // Get first positional argument if it exists + ann.args.first().and_then(|arg| { + match arg { + AnnotationArg::Positional(AnnotationValue::String(s)) => Some(s.clone()), + _ => None, + } + }) + }) +} + +fn get_encoding_annotation(annotations: &[Annotation], param_name: &str) -> Option { + annotations.iter() + .find(|ann| ann.name.name == "encoding") + .and_then(|ann| { + // First check for named argument matching param_name + for arg in &ann.args { + if let AnnotationArg::Named { name, value } = arg { + if name.name == param_name { + if let AnnotationValue::String(s) = value { + return Some(s.clone()); + } + } + } + } + + // Fall back to positional argument (applies to both input and output) + for arg in &ann.args { + if let AnnotationArg::Positional(AnnotationValue::String(s)) = arg { + return Some(s.clone()); + } + } + + None + }) +} + pub fn generate_lexicon(namespace: &str, lexicon: &Lexicon, workspace: &Workspace) -> Value { let usage_counts = analyze_type_usage(lexicon); @@ -319,10 +359,13 @@ fn generate_record_json(record: &Record, usage_counts: &HashMap, "properties": properties }); + // Check for @key annotation, default to "tid" + let key = get_annotation_string_value(&record.annotations, "key").unwrap_or_else(|| "tid".to_string()); + json!({ "type": "record", "description": extract_docs(&record.docs), - "key": "tid", + "key": key, "record": record_obj }) } @@ -358,11 +401,15 @@ fn generate_query_json(query: &Query, usage_counts: &HashMap, wor Value::Object(params_obj) }; + // Check for @encoding annotation (output only for queries), default to "application/json" + let output_encoding = get_encoding_annotation(&query.annotations, "output") + .unwrap_or_else(|| "application/json".to_string()); + let output = match &query.returns { ReturnType::None { .. } => None, ReturnType::Type(ty) => { let mut output_obj = Map::new(); - output_obj.insert("encoding".to_string(), json!("application/json")); + output_obj.insert("encoding".to_string(), json!(output_encoding)); output_obj.insert("schema".to_string(), generate_type_json(ty, usage_counts, workspace, current_namespace)); Some(Value::Object(output_obj)) } @@ -378,7 +425,7 @@ fn generate_query_json(query: &Query, usage_counts: &HashMap, wor } let mut output_obj = Map::new(); - output_obj.insert("encoding".to_string(), json!("application/json")); + output_obj.insert("encoding".to_string(), json!(output_encoding)); output_obj.insert("schema".to_string(), generate_type_json(success, usage_counts, workspace, current_namespace)); output_obj.insert("errors".to_string(), json!(error_defs)); Some(Value::Object(output_obj)) @@ -413,6 +460,10 @@ fn generate_procedure_json(procedure: &Procedure, usage_counts: &HashMap None, ReturnType::Type(ty) => { let mut output_obj = Map::new(); - output_obj.insert("encoding".to_string(), json!("application/json")); + output_obj.insert("encoding".to_string(), json!(output_encoding)); output_obj.insert("schema".to_string(), generate_type_json(ty, usage_counts, workspace, current_namespace)); Some(Value::Object(output_obj)) } @@ -447,7 +502,7 @@ fn generate_procedure_json(procedure: &Procedure, usage_counts: &HashMap