Enhance Type Builder API (#2177)

Fixes for #2050 & #2057
<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Enhance `TypeBuilder` API with new methods for property management and
class resetting, and update tests across Python, TypeScript, and Ruby.
> 
>   - **Behavior**:
> - Add `reset()` method to `TypeBuilder` and `ClassBuilder` in
`mod.rs`, `type_builder.rs`, and `type_builder.py` to clear properties
and classes.
> - Add `remove_property()` and `list_properties()` methods to
`ClassBuilder` in `mod.rs` and `type_builder.rs` for property
management.
> - Implement `get_type()` and `set_type()` in `ClassPropertyBuilder` in
`type_builder.rs`.
>   - **Tests**:
> - Update Python tests in `test_typebuilder.py` to cover `reset()`,
`remove_property()`, and `list_properties()`.
> - Update TypeScript tests in `dynamic-types.test.ts` to verify new
methods and dynamic type handling.
> - Add Ruby tests in `test_typebuilder.rb` for `reset()` and
`remove_property()` methods.
>   - **Misc**:
> - Update TypeScript and Ruby integration files to reflect new
`TypeBuilder` functionalities.
>     - Add documentation examples in `typebuilder.mdx` for new methods.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 51c4af9. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: antoniosarosi

.gitignore

@@ -79,6 +79,7 @@ $RECYCLE.BIN/
 .nfs*
 .node_repl_history
 .npm
+.pnpm-store/
 .nuxt
 .nyc_output
 .nyc_output/

engine/baml-runtime/src/type_builder/mod.rs

@@ -1,5 +1,6 @@
 use std::{
     fmt,
+    ops::Deref,
     sync::{Arc, Mutex},
 };
 
@@ -104,6 +105,18 @@ impl ClassBuilder {
         }
     }
 
+    // TODO: Figure out captured lifetime issue and return Iterator.
+    // Iterator that holds mutex lock seems tricky.
+    pub fn list_properties_key_value(&self) -> Vec<(String, ClassPropertyBuilder)> {
+        self.properties
+            .lock()
+            .unwrap()
+            .iter()
+            .map(|(name, prop)| (name.clone(), prop.lock().unwrap().deref().to_owned()))
+            .collect()
+    }
+
+    // TODO: Unify function above and this one (split because of CFFI).
     pub fn list_properties(&self) -> Vec<String> {
         let properties = self.properties.lock().unwrap();
         properties.keys().cloned().collect()
@@ -123,6 +136,15 @@ impl ClassBuilder {
             }))
         }))
     }
+
+    pub fn remove_property(&self, name: &str) {
+        let mut properties = self.properties.lock().unwrap();
+        properties.shift_remove(name);
+    }
+
+    pub fn reset(&self) {
+        self.properties.lock().unwrap().clear();
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -500,6 +522,14 @@ impl TypeBuilder {
         }
     }
 
+    pub fn reset(&self) {
+        self.classes.lock().unwrap().clear();
+        self.enums.lock().unwrap().clear();
+        self.type_aliases.lock().unwrap().clear();
+        self.recursive_type_aliases.lock().unwrap().clear();
+        self.recursive_classes.lock().unwrap().clear();
+    }
+
     pub fn upsert_class(&self, name: &str) -> Arc<Mutex<ClassBuilder>> {
         Arc::clone(
             self.classes

engine/generators/languages/python/src/generated_types.rs

@@ -180,7 +180,13 @@ mod type_builder {
     ///         return self._bldr.property(name).type(type)
     ///
     ///     def list_properties(self) -> typing.List[typing.Tuple[str, {{ class.class_property_type() }}]]:
-    ///         return [(name, self._bldr.property(name)) for name in self._properties]
+    ///         return self._bldr.list_properties()
+    ///
+    ///     def remove_property(self, name: str) -> None:
+    ///         self._bldr.remove_property(name)
+    ///
+    ///     def reset(self) -> None:
+    ///         self._bldr.reset()
     ///
     ///     {% else %}
     ///     def list_properties(self) -> typing.List[typing.Tuple[str, {{ class.class_property_type() }}]]:

engine/generators/languages/typescript/src/_templates/type_builder.ts.j2

@@ -37,6 +37,18 @@ export default class TypeBuilder {
         {% endfor %}
     }
 
+    reset(): void {
+        this.tb.reset();
+        // TODO: This should happen in Rust. Problem is, when we construct the
+        // typebuilder we instantiate class builders once and it seems to make
+        // a JS copy, bypassing the Rust side? In Python however, every time we
+        // access a class builder with @property, we get a new instance that
+        // wraps over the Rust type builder, so we only need to call tb.reset().
+        // In JS it's not possible unless we refactor the way class builders are
+        // accessed.
+        {% for cls in classes %}{% if cls.dynamic %} this.{{cls.name}}.reset(); {% endif %}{% endfor %}
+    }
+
     __tb() {
       return this.tb._tb();
     }

engine/language_client_python/python_src/baml_py/baml_py.pyi

@@ -275,6 +275,7 @@ class BamlSpan:
 
 class TypeBuilder:
     def __init__(self) -> None: ...
+    def reset(self) -> None: ...
     def enum(self, name: str) -> EnumBuilder: ...
     def class_(self, name: str) -> ClassBuilder: ...
     def string(self) -> FieldType: ...
@@ -470,6 +471,7 @@ class ClientRegistry:
 class FieldType:
     def list(self) -> FieldType: ...
     def optional(self) -> FieldType: ...
+    def __eq__(self, other: object) -> bool: ...
 
 class EnumBuilder:
     def value(self, name: str) -> EnumValueBuilder: ...
@@ -483,10 +485,14 @@ class EnumValueBuilder:
 
 class ClassBuilder:
     def field(self) -> FieldType: ...
+    def list_properties(self) -> List[Tuple[str, ClassPropertyBuilder]]: ...
     def property(self, name: str) -> ClassPropertyBuilder: ...
+    def remove_property(self, name: str) -> None: ...
+    def reset(self) -> None: ...
 
 class ClassPropertyBuilder:
     def type(self, field_type: FieldType) -> ClassPropertyBuilder: ...
+    def get_type(self) -> FieldType: ...
     def alias(self, alias: Optional[str]) -> ClassPropertyBuilder: ...
     def description(self, description: Optional[str]) -> ClassPropertyBuilder: ...
 

engine/language_client_python/python_src/baml_py/type_builder.py

@@ -19,6 +19,9 @@ def __init__(
         self.__tb = _TypeBuilder()
         self.__runtime = runtime
 
+    def reset(self):
+        self.__tb.reset()
+
     def __str__(self) -> str:
         """
         returns a comprehensive string representation of the typebuilder.
@@ -111,29 +114,22 @@ def add_baml(self, baml: str):
 class NewClassBuilder:
     def __init__(self, tb: _TypeBuilder, name: str):
         self.__bldr = tb.class_(name)
-        self.__properties: typing.Set[str] = set()
-        self.__props = NewClassProperties(self.__bldr, self.__properties)
 
     def type(self) -> FieldType:
         return self.__bldr.field()
 
     def list_properties(self) -> typing.List[typing.Tuple[str, "ClassPropertyBuilder"]]:
-        return [
-            (name, ClassPropertyBuilder(self.__bldr.property(name)))
-            for name in self.__properties
-        ]
+        return self.__bldr.list_properties()
+
+    def reset(self):
+        self.__bldr.reset()
+
+    def remove_property(self, name: str):
+        self.__bldr.remove_property(name)
 
     def add_property(self, name: str, type: FieldType) -> "ClassPropertyBuilder":
-        if name in self.__properties:
-            raise ValueError(f"Property {name} already exists.")
-        # BUG: we don't add to self.__properties here
-        # correct fix is to implement this logic in rust, not python
         return ClassPropertyBuilder(self.__bldr.property(name).type(type))
 
-    @property
-    def props(self) -> "NewClassProperties":
-        return self.__props
-
 
 class ClassPropertyBuilder:
     def __init__(self, bldr: _ClassPropertyBuilder):
@@ -153,17 +149,6 @@ def __init__(self, bldr: _ClassPropertyBuilder):
         self.__bldr = bldr
 
 
-class NewClassProperties:
-    def __init__(self, cls_bldr: ClassBuilder, properties: typing.Set[str]):
-        self.__bldr = cls_bldr
-        self.__properties = properties
-
-    def __getattr__(self, name: str) -> "ClassPropertyBuilder":
-        if name not in self.__properties:
-            raise AttributeError(f"Property {name} not found.")
-        return ClassPropertyBuilder(self.__bldr.property(name))
-
-
 class NewEnumBuilder:
     def __init__(self, tb: _TypeBuilder, name: str):
         self.__bldr = tb.enum(name)

engine/language_client_python/src/types/type_builder.rs

@@ -1,3 +1,5 @@
+use std::ops::Deref;
+
 use baml_runtime::type_builder::{self, WithMeta};
 use baml_types::{ir_type::UnionConstructor, BamlValue};
 use pyo3::{
@@ -37,6 +39,10 @@ impl TypeBuilder {
         type_builder::TypeBuilder::new().into()
     }
 
+    pub fn reset(&self) {
+        self.inner.reset();
+    }
+
     /// provides a detailed string representation of the typebuilder for python users.
     ///
     /// this method exposes the rust-implemented string formatting to python, ensuring
@@ -149,6 +155,10 @@ impl FieldType {
     pub fn optional(&self) -> FieldType {
         self.inner.lock().unwrap().clone().as_optional().into()
     }
+
+    pub fn __eq__(&self, other: &FieldType) -> bool {
+        self.inner.lock().unwrap().deref() == other.inner.lock().unwrap().deref()
+    }
 }
 
 #[pymethods]
@@ -213,6 +223,24 @@ impl ClassBuilder {
         baml_types::TypeIR::class(&self.name).into()
     }
 
+    pub fn list_properties(&self) -> Vec<(String, ClassPropertyBuilder)> {
+        self.inner
+            .lock()
+            .unwrap()
+            .list_properties_key_value()
+            .into_iter()
+            .map(|(name, prop_builder)| (name, prop_builder.into()))
+            .collect()
+    }
+
+    pub fn remove_property(&self, name: &str) {
+        self.inner.lock().unwrap().remove_property(name);
+    }
+
+    pub fn reset(&self) {
+        self.inner.lock().unwrap().reset();
+    }
+
     pub fn property(&self, name: &str) -> ClassPropertyBuilder {
         self.inner.lock().unwrap().upsert_property(name).into()
     }
@@ -228,6 +256,17 @@ impl ClassPropertyBuilder {
         self.inner.clone().into()
     }
 
+    pub fn get_type(&self) -> PyResult<FieldType> {
+        self.inner
+            .lock()
+            .unwrap()
+            .r#type()
+            .map(FieldType::from)
+            .ok_or_else(|| BamlError::from_anyhow(anyhow::anyhow!(
+                "attempted to read a property that has no defined type, this is likely an internal bug"
+            )))
+    }
+
     #[pyo3(signature = (alias = None))]
     pub fn alias(&self, alias: Option<&str>) -> Self {
         self.inner.lock().unwrap().with_meta(

engine/language_client_ruby/ext/ruby_ffi/src/types/type_builder.rs

@@ -40,6 +40,10 @@ impl TypeBuilder {
         }
     }
 
+    pub fn reset(&self) {
+        self.inner.reset();
+    }
+
     pub fn r#enum(&self, name: String) -> EnumBuilder {
         EnumBuilder {
             inner: self.inner.upsert_enum(name.as_str()),
@@ -141,6 +145,8 @@ impl TypeBuilder {
         cls.define_singleton_method("new", function!(TypeBuilder::new, 0))?;
         cls.define_method("to_s", method!(TypeBuilder::to_s, 0))?;
         cls.define_method("enum", method!(TypeBuilder::r#enum, 1))?;
+        // TODO: Not exposed, Ruby doesn't work right now.
+        // cls.define_method("reset", method!(TypeBuilder::reset, 0))?;
         // "class" is used by Kernel: https://ruby-doc.org/core-3.0.2/Kernel.html#method-i-class
         cls.define_method("class_", method!(TypeBuilder::class, 1))?;
         cls.define_method("list", method!(TypeBuilder::list, 1))?;
@@ -256,6 +262,14 @@ impl EnumValueBuilder {
 }
 
 impl ClassBuilder {
+    pub fn reset(&self) {
+        self.inner.lock().unwrap().reset();
+    }
+
+    pub fn remove_property(&self, name: String) {
+        self.inner.lock().unwrap().remove_property(name.as_str());
+    }
+
     pub fn field(&self) -> FieldType {
         baml_types::TypeIR::class(&self.name).into()
     }
@@ -273,6 +287,9 @@ impl ClassBuilder {
 
         cls.define_method("field", method!(ClassBuilder::field, 0))?;
         cls.define_method("property", method!(ClassBuilder::property, 1))?;
+        // TODO: Not exposed, Ruby doesn't work right now.
+        // cls.define_method("reset", method!(ClassBuilder::reset, 0))?;
+        // cls.define_method("remove_property", method!(ClassBuilder::remove_property, 1))?;
 
         Ok(())
     }

engine/language_client_typescript/native.d.ts

@@ -61,12 +61,16 @@ export declare class BamlVideo {
 }
 
 export declare class ClassBuilder {
+  listProperties(): unknown[]
+  removeProperty(name: string): void
+  reset(): void
   field(): FieldType
   property(name: string): ClassPropertyBuilder
 }
 
 export declare class ClassPropertyBuilder {
   setType(fieldType: FieldType): ClassPropertyBuilder
+  getType(): FieldType
   alias(alias?: string | undefined | null): ClassPropertyBuilder
   description(description?: string | undefined | null): ClassPropertyBuilder
 }
@@ -103,6 +107,7 @@ export declare class EnumValueBuilder {
 export declare class FieldType {
   list(): FieldType
   optional(): FieldType
+  equals(other: FieldType): boolean
 }
 
 export declare class FunctionLog {
@@ -215,6 +220,7 @@ export declare class TraceStats {
 
 export declare class TypeBuilder {
   constructor()
+  reset(): void
   getEnum(name: string): EnumBuilder
   getClass(name: string): ClassBuilder
   list(inner: FieldType): FieldType