Skip to content

Commit

Permalink
feat(python): improve TypeBuilder string representation to show prope…
Browse files Browse the repository at this point in the history 
…rties and values

- Add property details (aliases, descriptions) to class string representation

- Add value details (aliases, descriptions) to enum string representation

- Sort properties and values alphabetically

- Properly escape quotes in descriptions

- Store metadata in builders for consistent access
  • Loading branch information
@afyef
afyef committed Dec 18, 2024
1 parent 251085e commit 7311baa
Showing 1 changed file with 224 additions and 41 deletions.
265 changes: 224 additions & 41 deletions engine/language_client_python/python_src/baml_py/type_builder.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,17 @@
"""
╔══════════════════════════════════════════════════════════════════════════════╗
║ ║
║ BAML Type Builder ║
║ ║
║ a Python interface for creating and modifying BAML types dynamically. ║
║ this module provides a clean and intuitive way to: ║
║ - create and manage classes and enums ║
║ - add properties with aliases and descriptions ║
║ - add enum values with aliases and descriptions ║
║ ║
╚══════════════════════════════════════════════════════════════════════════════╝
"""

import typing
from .baml_py import (
ClassBuilder,
Expand All @@ -10,46 +24,116 @@


class TypeBuilder:
"""A builder for creating and modifying types at runtime.
hi, so this class provides a Python-friendly interface for creating and modifying BAML types dynamically.
it maintains sets of class and enum names, and provides a readable string representation for debugging.
the format goes like:
- Empty TypeBuilder: "TypeBuilder(empty)"
- With classes only: "TypeBuilder(Classes: ['Class1', 'Class2'])"
- With enums only: "TypeBuilder(Enums: ['Enum1', 'Enum2'])"
- With both: "TypeBuilder(Classes: ['Class1', 'Class2'], Enums: ['Enum1', 'Enum2'])"
a note from me: class and enum names are always sorted alphabetically for consistent output.
"""a builder for creating and modifying types at runtime.
┌─────────────────────────────────────────────────────────────────┐
│ the string representation format goes like this: │
│ ▸ Empty: TypeBuilder(empty) │
│ ▸ Classes: TypeBuilder(Classes: ['Class1', 'Class2']) │
│ ▸ Enums: TypeBuilder(Enums: ['Enum1', 'Enum2']) │
│ ▸ Both: TypeBuilder(Classes: [...], Enums: [...]) │
│ │
│ properties and values: │
│ ▸ With alias: name (alias='username') │
│ ▸ With description: name (desc='User name') │
│ ▸ With both: name (alias='x', desc='y') │
│ │
│ note: all names are sorted alphabetically for consistency │
└─────────────────────────────────────────────────────────────────┘
"""

def __init__(self, classes: typing.Set[str], enums: typing.Set[str]):
"""initializingt the TypeBuilder with optional predefined classes and enums.
args:
classes: Set of class names to initialize with
enums: Set of enum names to initialize with
"""initialize the TypeBuilder with optional predefined classes and enums.
┌────────────────────────────────────────────────────────────┐
│ internal State: │
│ ▸ __classes: Set of class names │
│ ▸ __enums: Set of enum names │
│ ▸ __tb: Rust TypeBuilder implementation │
│ ▸ __class_builders: Dict of class name -> builder │
│ ▸ __enum_builders: Dict of enum name -> builder │
└────────────────────────────────────────────────────────────┘
"""
self.__classes = classes
self.__enums = enums
self.__tb = _TypeBuilder()
self.__class_builders: typing.Dict[str, NewClassBuilder] = {}
self.__enum_builders: typing.Dict[str, NewEnumBuilder] = {}

def __str__(self) -> str:
"""here, we create a human-readable string representation of the TypeBuilder as required.
what this actually returns:
- a string showing all currently defined classes and enums in a readable format.
- classes and enums are sorted alphabetically for consistent output.
- returns "TypeBuilder(empty)" if no types are defined.
"""a human-readable string representation of the typebuilder.
+-----------------------------------------------------------
| string representation features: |
| ▸ shows all defined classes and enums |
| ▸ includes property types, aliases, and descriptions |
| ▸ includes enum values with aliases and descriptions |
| ▸ sorts everything alphabetically for consistency |
| ▸ uses 'typebuilder(empty)' when no types are defined |
| |
| example: |
| typebuilder( |
| classes: ['user { name (alias='x', desc='y') }'], |
| enums: ['status { active (desc='active state') }'] |
| ) |
+-----------------------------------------------------------+
"""
parts = []
# add sorted class names if exist
# add sorted class names with their properties if they exist
if self.__classes:
parts.append(f"Classes: {sorted(self.__classes)}")
# add sorted enum names if exist
class_details = []
for class_name in sorted(self.__classes):
class_builder = self.__class_builders[class_name]
props = []
# Get all properties and sort them
property_list = class_builder.list_properties()
# Sort by property name, but put properties with aliases first
property_list.sort(key=lambda x: (
0 if x[1].get_alias() else 1, # Properties with aliases first
x[0] # Then by name
))
for prop_name, prop_builder in property_list:
prop_str = prop_name
alias = prop_builder.get_alias()
desc = prop_builder.get_description()
if alias or desc:
extras = []
if alias:
extras.append(f"alias='{alias}'")
if desc:
extras.append(f"desc='{desc.replace("'", "\\'")}'")
prop_str += f" ({', '.join(extras)})"
props.append(prop_str)
class_str = class_name
if props:
class_str += f" {{ {', '.join(props)} }}"
class_details.append(class_str)
parts.append(f"Classes: ['{', '.join(class_details)}']")

# add sorted enum names with their values if they exist
if self.__enums:
parts.append(f"Enums: {sorted(self.__enums)}")
enum_details = []
for enum_name in sorted(self.__enums):
enum_builder = self.__enum_builders[enum_name]
values = []
for value_name, value_builder in enum_builder.list_values():
value_str = value_name
alias = value_builder.get_alias()
desc = value_builder.get_description()
if alias or desc:
extras = []
if alias:
extras.append(f"alias='{alias}'")
if desc:
extras.append(f"desc='{desc.replace("'", "\\'")}'")
value_str += f" ({', '.join(extras)})"
values.append(value_str)
enum_str = enum_name
if values:
enum_str += f" {{ {', '.join(values)} }}"
enum_details.append(enum_str)
parts.append(f"Enums: ['{', '.join(enum_details)}']")

# return special format for empty TypeBuilder
if not parts:
return "TypeBuilder(empty)"
Expand Down Expand Up @@ -103,56 +187,99 @@ def add_class(self, name: str) -> "NewClassBuilder":
if name in self.__enums:
raise ValueError(f"Enum with name {name} already exists.")
self.__classes.add(name)
return NewClassBuilder(self._tb, name)
builder = NewClassBuilder(self._tb, name)
self.__class_builders[name] = builder
return builder

def add_enum(self, name: str) -> "NewEnumBuilder":
if name in self.__classes:
raise ValueError(f"Class with name {name} already exists.")
if name in self.__enums:
raise ValueError(f"Enum with name {name} already exists.")
self.__enums.add(name)
return NewEnumBuilder(self._tb, name)
builder = NewEnumBuilder(self._tb, name)
self.__enum_builders[name] = builder
return builder


class NewClassBuilder:
"""builder for creating and modifying BAML classes.
┌─────────────────────────────────────────────────────────────┐
│ features: │
│ ▸ add properties with types │
│ ▸ set aliases for better field mapping │
│ ▸ add descriptions for documentation │
│ ▸ track properties for consistent access │
│ │
│ example: │
│ user = builder.add_class("user") │
│ user.add_property("name", tb.string()) │
│ .alias("username") │
│ .description("the user's full name") │
└─────────────────────────────────────────────────────────────┘
"""

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)
self.__property_builders: typing.Dict[str, ClassPropertyBuilder] = {}

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 [(name, self.__property_builders[name]) for name in sorted(self.__properties)]

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))
self.__properties.add(name)
builder = ClassPropertyBuilder(self.__bldr.property(name).type(type))
self.__property_builders[name] = builder
return builder

@property
def props(self) -> "NewClassProperties":
return self.__props


class ClassPropertyBuilder:
"""builder for configuring class properties.
┌─────────────────────────────────────────────────────────────┐
│ features: │
│ ▸ set aliases for field mapping │
│ ▸ add descriptions for documentation │
│ ▸ store metadata for string representation │
│ │
│ example: │
│ property.alias("username") │
│ .description("the user's full name") │
└─────────────────────────────────────────────────────────────┘
"""

def __init__(self, bldr: _ClassPropertyBuilder):
self.__bldr = bldr
self.__alias = None
self.__description = None

def alias(self, alias: typing.Optional[str]):
self.__alias = alias
self.__bldr.alias(alias)
return self

def description(self, description: typing.Optional[str]):
self.__description = description
self.__bldr.description(description)
return self

def get_alias(self) -> typing.Optional[str]:
return self.__alias

def get_description(self) -> typing.Optional[str]:
return self.__description


class NewClassProperties:
def __init__(self, cls_bldr: ClassBuilder, properties: typing.Set[str]):
Expand All @@ -166,10 +293,28 @@ def __getattr__(self, name: str) -> "ClassPropertyBuilder":


class NewEnumBuilder:
"""builder for creating and modifying BAML enums.
┌─────────────────────────────────────────────────────────────┐
│ features: │
│ ▸ add enum values │
│ ▸ set aliases for better value mapping │
│ ▸ add descriptions for documentation │
│ ▸ track values for consistent access │
│ │
│ example: │
│ status = builder.add_enum("status") │
│ status.add_value("active") │
│ .alias("active") │
│ .description("user is active") │
└─────────────────────────────────────────────────────────────┘
"""

def __init__(self, tb: _TypeBuilder, name: str):
self.__bldr = tb.enum(name)
self.__values: typing.Set[str] = set()
self.__vals = NewEnumValues(self.__bldr, self.__values)
self.__value_builders: typing.Dict[str, "EnumValueBuilderWrapper"] = {}

def type(self) -> FieldType:
return self.__bldr.field()
Expand All @@ -178,15 +323,53 @@ def type(self) -> FieldType:
def values(self) -> "NewEnumValues":
return self.__vals

def list_values(self) -> typing.List[typing.Tuple[str, EnumValueBuilder]]:
return [(name, self.__bldr.value(name)) for name in self.__values]
def list_values(self) -> typing.List[typing.Tuple[str, "EnumValueBuilderWrapper"]]:
return [(name, self.__value_builders[name]) for name in sorted(self.__values)]

def add_value(self, name: str) -> "EnumValueBuilder":
def add_value(self, name: str) -> "EnumValueBuilderWrapper":
if name in self.__values:
raise ValueError(f"Value {name} already exists.")
self.__values.add(name)
# NOTE(sam): why is this inconsistent between classes and enums?
return self.__bldr.value(name)
builder = EnumValueBuilderWrapper(self.__bldr.value(name))
self.__value_builders[name] = builder
return builder


class EnumValueBuilderWrapper:
"""builder for configuring enum values.
┌─────────────────────────────────────────────────────────────┐
│ features: │
│ ▸ set aliases for value mapping │
│ ▸ add descriptions for documentation │
│ ▸ store metadata for string representation │
│ │
│ example: │
│ value.alias("active") │
│ .description("user is active") │
└─────────────────────────────────────────────────────────────┘
"""

def __init__(self, bldr: EnumValueBuilder):
self.__bldr = bldr
self.__alias = None
self.__description = None

def alias(self, alias: typing.Optional[str]):
self.__alias = alias
self.__bldr.alias(alias)
return self

def description(self, description: typing.Optional[str]):
self.__description = description
self.__bldr.description(description)
return self

def get_alias(self) -> typing.Optional[str]:
return self.__alias

def get_description(self) -> typing.Optional[str]:
return self.__description


class NewEnumValues:
Expand Down

0 comments on commit 7311baa

Please sign in to comment.