use super::optional_renderers::{push_key_field_name, push_key_index_suffix};
use super::renderers::quoted_key_literal;
use super::types::{PathSegment, PythonTypedDictMap};
use std::collections::HashSet;
pub(super) fn render_python_with_optionals(
segments: &[PathSegment],
result_var: &str,
optional_fields: &HashSet<String>,
typeddict_map: &PythonTypedDictMap,
) -> String {
let last_index = segments.len().saturating_sub(1);
let mut crossings: Vec<usize> = Vec::new();
let mut path_so_far = String::new();
for (index, segment) in segments.iter().enumerate() {
push_key_field_name(&mut path_so_far, segment);
if index != last_index && optional_fields.contains(&path_so_far) {
crossings.push(index);
}
push_key_index_suffix(&mut path_so_far, segment);
}
let mut expression = render_python_accessor(segments, result_var, typeddict_map);
for &index in crossings.iter().rev() {
let condition = render_python_accessor(&field_only_prefix(segments, index), result_var, typeddict_map);
expression = format!("({expression} if {condition} else None)");
}
expression
}
pub(super) fn render_python_accessor(segments: &[PathSegment], result_var: &str, map: &PythonTypedDictMap) -> String {
let mut out = result_var.to_string();
let mut current_type = map.root_type.clone();
for seg in segments {
match seg {
PathSegment::Field(f) => {
push_field_access(&mut out, f, current_type.as_deref(), map);
current_type = map.advance(current_type.as_deref(), f);
}
PathSegment::ArrayField { name, index } => {
push_field_access(&mut out, name, current_type.as_deref(), map);
out.push_str(&format!("[{index}]"));
current_type = map.advance(current_type.as_deref(), name);
}
PathSegment::MapAccess { field, key } => {
push_field_access(&mut out, field, current_type.as_deref(), map);
if key.chars().all(|c| c.is_ascii_digit()) {
let idx: usize = key.parse().unwrap_or(0);
out.push_str(&format!("[{idx}]"));
} else {
out.push_str(&format!(".get({})", quoted_key_literal(key)));
}
}
PathSegment::Length => {
let current = std::mem::take(&mut out);
out = format!("len({current})");
}
}
}
out
}
fn push_field_access(out: &mut String, field: &str, owner_type: Option<&str>, map: &PythonTypedDictMap) {
if map.is_typeddict(owner_type) {
out.push_str(&format!("[{}]", quoted_key_literal(field)));
} else {
out.push('.');
out.push_str(field);
}
}
fn field_only_prefix(segments: &[PathSegment], index: usize) -> Vec<PathSegment> {
let mut prefix: Vec<PathSegment> = segments[..index].to_vec();
prefix.push(match &segments[index] {
PathSegment::ArrayField { name, .. } => PathSegment::Field(name.clone()),
PathSegment::MapAccess { field, .. } => PathSegment::Field(field.clone()),
other => other.clone(),
});
prefix
}
#[cfg(test)]
mod tests {
use super::super::parse::parse_path;
use super::*;
fn typeddict_map(
typeddict_types: &[&str],
field_types: &[(&str, &str, &str)],
root_type: &str,
) -> PythonTypedDictMap {
let mut map = PythonTypedDictMap {
typeddict_types: typeddict_types.iter().map(|s| s.to_string()).collect(),
root_type: Some(root_type.to_string()),
..Default::default()
};
for (owner, field, target) in field_types {
map.field_types
.entry(owner.to_string())
.or_default()
.insert(field.to_string(), target.to_string());
}
map
}
#[test]
fn a_scalar_field_on_a_typeddict_result_is_subscripted() {
let map = typeddict_map(&["ApiResult"], &[], "ApiResult");
let segments = parse_path("status_code");
assert_eq!(
render_python_accessor(&segments, "result", &map),
r#"result["status_code"]"#
);
}
#[test]
fn a_scalar_field_on_a_non_typeddict_result_stays_attribute_access() {
let map = PythonTypedDictMap::default();
let segments = parse_path("status_code");
assert_eq!(render_python_accessor(&segments, "result", &map), "result.status_code");
}
#[test]
fn a_typeddict_result_with_an_optional_field_narrows_via_subscript() {
let map = typeddict_map(&["ApiResult"], &[], "ApiResult");
let optional: HashSet<String> = ["markdown".to_string()].into_iter().collect();
let segments = parse_path("markdown");
assert_eq!(
render_python_with_optionals(&segments, "result", &optional, &map),
r#"result["markdown"]"#,
"a crossing at the LAST segment needs no ternary guard"
);
}
#[test]
fn a_typeddict_result_with_an_optional_nested_typeddict_field_narrows_before_descending() {
let map = typeddict_map(
&["ApiResult", "Markdown"],
&[("ApiResult", "markdown", "Markdown")],
"ApiResult",
);
let optional: HashSet<String> = ["markdown".to_string()].into_iter().collect();
let segments = parse_path("markdown.content");
assert_eq!(
render_python_with_optionals(&segments, "result", &optional, &map),
r#"(result["markdown"]["content"] if result["markdown"] else None)"#
);
}
#[test]
fn descending_from_a_typeddict_into_a_non_typeddict_nested_type_switches_to_attribute_access() {
let map = typeddict_map(&["ApiResult"], &[("ApiResult", "metadata", "Metadata")], "ApiResult");
let segments = parse_path("metadata.title");
assert_eq!(
render_python_accessor(&segments, "result", &map),
r#"result["metadata"].title"#
);
}
#[test]
fn an_array_field_on_a_typeddict_result_subscripts_the_field_then_indexes_the_list() {
let map = typeddict_map(&["ApiResult"], &[], "ApiResult");
let segments = parse_path("pages[0]");
assert_eq!(
render_python_accessor(&segments, "result", &map),
r#"result["pages"][0]"#
);
}
#[test]
fn indexing_past_an_optional_field_narrows_it_first() {
let map = PythonTypedDictMap::default();
let optional: HashSet<String> = ["choices[0].message.tool_calls".to_string()].into_iter().collect();
let segments = parse_path("choices[0].message.tool_calls[0].function.name");
assert_eq!(
render_python_with_optionals(&segments, "result", &optional, &map),
"(result.choices[0].message.tool_calls[0].function.name if result.choices[0].message.tool_calls else None)"
);
}
}