Skip to content

Commit

Permalink
Autogenerated partial updates APIs for Python (#8671)
Browse files Browse the repository at this point in the history 
There really is not much to be said that wouldn't be made clearer by
just looking at the code.

Specifically:
* Look at the codegen changes.
* Now look at the changes to one of the generated archetypes (probably
`Transform3D`).
* Now look especially at the changes to `transform3d_ext`. That one's
important.
* Now look at how the different "partial updates" snippets evolved.

Overall I hate everything here, literally every single bit -- but that's
the best I managed to get out of the status quo without ridiculous
amounts of efforts.

* DNM: requires #8690 
* Closes #8582

---------

Co-authored-by: Antoine Beyeler <antoine@rerun.io>
  • Loading branch information
@teh-cmc @abey79
teh-cmc and abey79 authored Jan 15, 2025
1 parent ac80fbc commit b19ec4f
Show file tree
Hide file tree
Showing 94 changed files with 5,570 additions and 1,289 deletions.
218 changes: 156 additions & 62 deletions crates/build/re_types_builder/src/codegen/python/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ use crate::{
objects::ObjectClass,
ArrowRegistry, CodeGenerator, Docs, ElementType, GeneratedFiles, Object, ObjectField,
ObjectKind, Objects, Reporter, Type, ATTR_PYTHON_ALIASES, ATTR_PYTHON_ARRAY_ALIASES,
ATTR_RERUN_LOG_MISSING_AS_EMPTY,
};

use self::views::code_for_view;
Expand Down Expand Up @@ -592,11 +591,7 @@ fn code_for_struct(
} else if *kind == ObjectKind::Archetype {
// Archetypes use the ComponentBatch constructor for their fields
let (typ_unwrapped, _) = quote_field_type_from_field(objects, field, true);
if field.is_nullable && !obj.attrs.has(ATTR_RERUN_LOG_MISSING_AS_EMPTY) {
format!("converter={typ_unwrapped}Batch._optional, # type: ignore[misc]\n")
} else {
format!("converter={typ_unwrapped}Batch._required, # type: ignore[misc]\n")
}
format!("converter={typ_unwrapped}Batch._converter, # type: ignore[misc]\n")
} else if !default_converter.is_empty() {
code.push_indented(0, &converter_function, 1);
format!("converter={default_converter}")
Expand Down Expand Up @@ -699,6 +694,7 @@ fn code_for_struct(

if obj.kind == ObjectKind::Archetype {
code.push_indented(1, quote_clear_methods(obj), 2);
code.push_indented(1, quote_partial_update_methods(reporter, obj, objects), 2);
}

if obj.is_delegating_component() {
Expand Down Expand Up @@ -739,10 +735,7 @@ fn code_for_struct(
};

let metadata = if *kind == ObjectKind::Archetype {
format!(
"\nmetadata={{'component': '{}'}}, ",
if *is_nullable { "optional" } else { "required" }
)
"\nmetadata={'component': True}, ".to_owned()
} else {
String::new()
};
Expand All @@ -756,7 +749,7 @@ fn code_for_struct(
String::new()
};
// Note: mypy gets confused using staticmethods for field-converters
let typ = if !*is_nullable {
let typ = if !obj.is_archetype() && !*is_nullable {
format!("{typ} = field(\n{metadata}{converter}{type_ignore}\n)")
} else {
format!(
Expand Down Expand Up @@ -2336,60 +2329,57 @@ fn quote_init_parameter_from_field(
}
}

fn quote_init_method(
reporter: &Reporter,
obj: &Object,
ext_class: &ExtensionClass,
objects: &Objects,
) -> String {
fn compute_init_parameters(obj: &Object, objects: &Objects) -> Vec<String> {
// If the type is fully transparent (single non-nullable field and not an archetype),
// we have to use the "{obj.name}Like" type directly since the type of the field itself might be too narrow.
// -> Whatever type aliases there are for this type, we need to pick them up.
let parameters: Vec<_> =
if obj.kind != ObjectKind::Archetype && obj.fields.len() == 1 && !obj.fields[0].is_nullable
{
vec![format!(
"{}: {}",
obj.fields[0].name,
quote_parameter_type_alias(&obj.fqname, &obj.fqname, objects, false)
)]
} else if obj.is_union() {
vec![format!(
"inner: {} | None = None",
quote_parameter_type_alias(&obj.fqname, &obj.fqname, objects, false)
)]
} else {
let required = obj
.fields
.iter()
.filter(|field| !field.is_nullable)
.map(|field| quote_init_parameter_from_field(field, objects, &obj.fqname))
.collect_vec();

let optional = obj
.fields
.iter()
.filter(|field| field.is_nullable)
.map(|field| quote_init_parameter_from_field(field, objects, &obj.fqname))
.collect_vec();

if optional.is_empty() {
required
} else if obj.kind == ObjectKind::Archetype {
// Force kw-args for all optional arguments:
required
.into_iter()
.chain(std::iter::once("*".to_owned()))
.chain(optional)
.collect()
} else {
required.into_iter().chain(optional).collect()
}
};
if obj.kind != ObjectKind::Archetype && obj.fields.len() == 1 && !obj.fields[0].is_nullable {
vec![format!(
"{}: {}",
obj.fields[0].name,
quote_parameter_type_alias(&obj.fqname, &obj.fqname, objects, false)
)]
} else if obj.is_union() {
vec![format!(
"inner: {} | None = None",
quote_parameter_type_alias(&obj.fqname, &obj.fqname, objects, false)
)]
} else {
let required = obj
.fields
.iter()
.filter(|field| !field.is_nullable)
.map(|field| quote_init_parameter_from_field(field, objects, &obj.fqname))
.collect_vec();

let head = format!("def __init__(self: Any, {}):", parameters.join(", "));
let optional = obj
.fields
.iter()
.filter(|field| field.is_nullable)
.map(|field| quote_init_parameter_from_field(field, objects, &obj.fqname))
.collect_vec();

if optional.is_empty() {
required
} else if obj.kind == ObjectKind::Archetype {
// Force kw-args for all optional arguments:
required
.into_iter()
.chain(std::iter::once("*".to_owned()))
.chain(optional)
.collect()
} else {
required.into_iter().chain(optional).collect()
}
}
}

let parameter_docs = if obj.is_union() {
fn compute_init_parameter_docs(
reporter: &Reporter,
obj: &Object,
objects: &Objects,
) -> Vec<String> {
if obj.is_union() {
Vec::new()
} else {
obj.fields
Expand All @@ -2414,7 +2404,19 @@ fn quote_init_method(
}
})
.collect::<Vec<_>>()
};
}
}

fn quote_init_method(
reporter: &Reporter,
obj: &Object,
ext_class: &ExtensionClass,
objects: &Objects,
) -> String {
let parameters = compute_init_parameters(obj, objects);
let head = format!("def __init__(self: Any, {}):", parameters.join(", "));

let parameter_docs = compute_init_parameter_docs(reporter, obj, objects);
let mut doc_string_lines = vec![format!(
"Create a new instance of the {} {}.",
obj.name,
Expand Down Expand Up @@ -2474,7 +2476,7 @@ fn quote_clear_methods(obj: &Object) -> String {
let param_nones = obj
.fields
.iter()
.map(|field| format!("{} = None, # type: ignore[arg-type]", field.name))
.map(|field| format!("{} = None,", field.name))
.join("\n ");

let classname = &obj.name;
Expand All @@ -2497,6 +2499,98 @@ fn quote_clear_methods(obj: &Object) -> String {
))
}

fn quote_partial_update_methods(reporter: &Reporter, obj: &Object, objects: &Objects) -> String {
let name = &obj.name;

let parameters = obj
.fields
.iter()
.map(|field| {
let mut field = field.clone();
field.is_nullable = true;
quote_init_parameter_from_field(&field, objects, &obj.fqname)
})
.collect_vec()
.join(", ");

let kwargs = obj
.fields
.iter()
.map(|field| {
let field_name = field.snake_case_name();
format!("'{field_name}': {field_name}")
})
.collect_vec()
.join(", ");

let parameter_docs = compute_init_parameter_docs(reporter, obj, objects);
let mut doc_string_lines = vec![format!("Update only some specific fields of a `{name}`.")];
if !parameter_docs.is_empty() {
doc_string_lines.push("\n".to_owned());
doc_string_lines.push("Parameters".to_owned());
doc_string_lines.push("----------".to_owned());
doc_string_lines.push("clear:".to_owned());
doc_string_lines
.push(" If true, all unspecified fields will be explicitly cleared.".to_owned());
for doc in parameter_docs {
doc_string_lines.push(doc);
}
};
let doc_block = quote_doc_lines(doc_string_lines)
.lines()
.map(|line| format!(" {line}"))
.collect_vec()
.join("\n");

let field_clears = obj
.fields
.iter()
.map(|field| {
let field_name = field.snake_case_name();
format!("{field_name}=[],")
})
.collect_vec()
.join("\n ");
let field_clears = indent::indent_by(4, field_clears);

unindent(&format!(
r#"
@classmethod
def update_fields(
cls,
*,
clear: bool = False,
{parameters},
) -> {name}:
{doc_block}
inst = cls.__new__(cls)
with catch_and_log_exceptions(context=cls.__name__):
kwargs = {{
{kwargs},
}}
if clear:
kwargs = {{k: v if v is not None else [] for k, v in kwargs.items()}} # type: ignore[misc]
inst.__attrs_init__(**kwargs)
return inst
inst.__attrs_clear__()
return inst
@classmethod
def clear_fields(cls) -> {name}:
"""Clear all the fields of a `{name}`."""
inst = cls.__new__(cls)
inst.__attrs_init__(
{field_clears}
)
return inst
"#
))
}

// --- Arrow registry code generators ---
use arrow2::datatypes::{DataType, Field, UnionMode};

Expand Down
Loading

0 comments on commit b19ec4f

Please sign in to comment.