type_bridge.generator

generator

TypeBridge code generator - Generate Python models from TypeDB schemas.

This module provides tools to parse TypeDB TQL schema files and generate corresponding type-bridge Python model classes.

Example usage:

from type_bridge.generator import generate_models

# Generate from a schema file
generate_models("schema.tql", "myapp/models/")

# Or from schema text
schema_text = '''
define
entity person,
    owns name @key;
attribute name, value string;
'''
generate_models(schema_text, "myapp/models/")

The generated package structure:

myapp/models/
├── __init__.py      # Package exports, ATTRIBUTES/ENTITIES/RELATIONS lists
├── attributes.py    # Attribute class definitions
├── entities.py      # Entity class definitions
├── relations.py     # Relation class definitions
└── schema.tql       # Copy of the source schema

BaseClassConfig dataclass

BaseClassConfig(source_entity, base_name, inherited_attrs=list(), extra_fields=dict(), field_syncs=list(), validators=list(), create_field_overrides=dict())

Configuration for a custom base class in the DTO hierarchy.

When an entity inherits from source_entity in the schema, its DTOs will inherit from the custom base class instead of the default.

Attributes:

Name	Type	Description
source_entity	str	The schema entity name that triggers this base class (e.g., "artifact"). Entities that inherit from this entity will use the custom base.
base_name	str	The base class name prefix (e.g., "BaseArtifact"). Will generate BaseArtifactOut, BaseArtifactCreate, BaseArtifactPatch.
inherited_attrs	list[str]	Attribute names that are defined in the base class. These will be skipped when rendering child DTOs to avoid duplication.
extra_fields	dict[str, str]	Additional fields to add to the base class. Keys are field names, values are type annotations as strings.
field_syncs	list[FieldSyncConfig]	Field sync configurations for this base class.
validators	list[str]	List of validator names to apply to fields in this base.

Example

BaseClassConfig( source_entity="artifact", base_name="BaseArtifact", inherited_attrs=["display_id", "name", "description", "status"], extra_fields={"version": "int | None = None"}, field_syncs=[FieldSyncConfig("description", "content")], )

CompositeEntityConfig dataclass

CompositeEntityConfig(name, base_entity=None, include_entities=list(), exclude_entities=list(), common_fields=list(), field_syncs=list(), extra_fields=dict(), extra_fields_out=dict(), skip_variants=set(), id_field_name='id', type_field_name='type', type_enum_from_registry=True)

Configuration for a composite/flat DTO that merges multiple entity types.

Composite DTOs provide a polymorphic API where a single DTO class handles multiple entity types. This is useful for graph APIs where you want one endpoint to handle all node types.

Attributes:

Name	Type	Description
name	str	Name prefix for the composite (e.g., "GraphNode" generates GraphNodeOut, GraphNodeCreate, GraphNodePatch)
base_entity	str | None	Schema entity name - all entities inheriting from this will be included in the composite (e.g., "artifact")
include_entities	list[str]	Explicit list of entity names to include. Use this OR base_entity, not both.
exclude_entities	list[str]	Entity names to exclude (useful with base_entity)
common_fields	list[CompositeFieldConfig]	Fields that appear on all variants with specified types. These are required (or have defaults) in Create, present in Out/Patch.
field_syncs	list[FieldSyncConfig]	Field sync configurations for the composite
extra_fields	dict[str, str]	Additional fields not in schema (e.g., computed fields). Applied to all variants (Out, Create, Patch) unless overridden.
extra_fields_out	dict[str, str]	Per-variant override for extra_fields on Out DTOs. If a field name appears here, the Out variant uses this annotation instead of the one from extra_fields. Useful when Out fields have different requiredness (e.g., id required on Out but optional on Create).
skip_variants	set[str]	Set of variant names to skip generating (e.g., {"out"} to skip the Out class). Valid values: "out", "create", "patch". The type enum is always generated regardless of this setting.
id_field_name	str	Name of the ID field (default: "id")
type_field_name	str	Name of the type discriminator field (default: "type")
type_enum_from_registry	bool	If True, generate dynamic enum from registry

Example

CompositeEntityConfig( name="GraphNode", base_entity="artifact", exclude_entities=["design_aspect"], common_fields=[ CompositeFieldConfig("name", "str"), CompositeFieldConfig("content", "str", "''"), CompositeFieldConfig("status", "str", "'proposed'"), ], field_syncs=[FieldSyncConfig("description", "content")], extra_fields={"version": "int | None = None"}, )

get_included_entities

get_included_entities(schema_entities)

Get list of entity names included in this composite.

Parameters:

Name	Type	Description	Default
schema_entities	dict[str, Any]	Dict of entity name -> EntitySpec from parsed schema	required

Returns:

Type	Description
list[str]	List of entity names to include in the composite

Source code in type_bridge/generator/dto_config.py

def get_included_entities(self, schema_entities: dict[str, Any]) -> list[str]:
    """Get list of entity names included in this composite.

    Args:
        schema_entities: Dict of entity name -> EntitySpec from parsed schema

    Returns:
        List of entity names to include in the composite
    """
    if self.include_entities:
        # Explicit list provided
        return [e for e in self.include_entities if e not in self.exclude_entities]

    if self.base_entity:
        # Find all entities that inherit from base_entity
        included = []
        for name, spec in schema_entities.items():
            if name in self.exclude_entities:
                continue
            if spec.abstract:
                continue
            # Check if this entity inherits from base_entity
            current = name
            while current:
                if current == self.base_entity:
                    included.append(name)
                    break
                parent_spec = schema_entities.get(current)
                if parent_spec and parent_spec.parent:
                    current = parent_spec.parent
                else:
                    break
        return included

    return []

CompositeFieldConfig dataclass

CompositeFieldConfig(name, type_annotation, default=None, description=None)

Configuration for a field in a composite entity DTO.

Attributes:

Name	Type	Description
name	str	Field name in Python (snake_case)
type_annotation	str	Python type annotation as string (e.g., "str", "int | None")
default	str | None	Default value as string (e.g., "None", "'proposed'", "''") If None, field is required in Create DTOs.
description	str | None	Optional field description for documentation

Example

CompositeFieldConfig( name="status", type_annotation="str", default="'proposed'", description="Lifecycle status" )

DTOConfig dataclass

DTOConfig(base_classes=list(), validators=list(), preamble=None, entity_union_name='Entity', relation_union_name='Relation', exclude_entities=list(), iid_field_name='iid', skip_relation_output=False, relation_create_base_class=None, relation_preamble=None, composite_entities=list(), strict_out_models=False, entity_field_overrides=list())

Configuration for API DTO generation.

This configuration controls how Pydantic DTOs are generated from a
TypeDB schema. It allows customization of base classes, validators,
field syncs, and naming conventions.

Attributes:
    base_classes: Custom base class configurations for entity hierarchies.
    validators: Custom validator type configurations.
    preamble: Custom Python code to inject at the top of the file
        (after imports). Use for complex validators or helper functions.
    entity_union_name: Name for the entity union type (default: "Entity").
    relation_union_name: Name for the relation union type (default: "Relation").
    exclude_entities: Entity names to exclude from generation (e.g., internal
        utility entities like "display_id_counter").
    iid_field_name: Field name for TypeDB IID in output DTOs (default: "iid").
        Set to "id" for more conventional REST API naming.
    skip_relation_output: If True, don't generate XxxOut classes for relations.
        Use with relation_preamble to define custom output classes.
    relation_create_base_class: Name of custom base class for relation Create
        DTOs (must be defined in preamble). If set, relation Create classes
        won't generate role-specific fields, only owned attributes.
    relation_preamble: Custom Python code to inject in the relation section.
        Use for custom relation output classes like GraphEdgeOut.
    composite_entities: Composite entity configurations for creating flat/merged
        DTOs that handle multiple entity types in one class.
    strict_out_models: If True, fields with @key or min>=1 cardinality will be
        required (not Optional) in Out DTOs. Default False for safety.

Example:
    config = DTOConfig(
        base_classes=[
            BaseClassConfig(
                source_entity="artifact",
                base_name="BaseArtifact",
                inherited_attrs=["display_id", "name"],
            ),
        ],
        validators=[
            ValidatorConfig(name="DisplayId", pattern=r"^[A-Z]{1,10}-\d+$"),
        ],
        exclude_entities=["display_id_counter", "schema_status"],
        iid_field_name="id",
        preamble='''

Custom node ID validator

def _validate_node_id(value: str) -> str: # Complex validation logic here return value

NodeId = Annotated[str, AfterValidator(_validate_node_id)] ''', )

get_base_class_for_entity

get_base_class_for_entity(entity_name, schema_entities)

Find the base class config for an entity based on inheritance.

Walks up the inheritance chain to find if this entity or any ancestor matches a configured base class.

Parameters:

Name	Type	Description	Default
entity_name	str	The entity name to check	required
schema_entities	dict[str, Any]	Dict of entity name -> EntitySpec from parsed schema	required

Returns:

Type	Description
BaseClassConfig | None	The matching BaseClassConfig, or None if no match

Source code in type_bridge/generator/dto_config.py

def get_base_class_for_entity(
    self, entity_name: str, schema_entities: dict[str, Any]
) -> BaseClassConfig | None:
    """Find the base class config for an entity based on inheritance.

    Walks up the inheritance chain to find if this entity or any ancestor
    matches a configured base class.

    Args:
        entity_name: The entity name to check
        schema_entities: Dict of entity name -> EntitySpec from parsed schema

    Returns:
        The matching BaseClassConfig, or None if no match
    """
    if not self.base_classes:
        return None

    # Build a set of source entities for quick lookup
    source_entities = {bc.source_entity: bc for bc in self.base_classes}

    # Walk up the inheritance chain
    current = entity_name
    while current:
        if current in source_entities:
            return source_entities[current]
        # Get parent
        entity_spec = schema_entities.get(current)
        if entity_spec and entity_spec.parent:
            current = entity_spec.parent
        else:
            break

    return None

EntityFieldOverride dataclass

EntityFieldOverride(entity, field, variant, required=None, default=None)

Per-entity, per-variant field override.

Allows targeting a specific entity + field + variant combination.

Attributes:

Name	Type	Description
entity	str	TypeDB entity name (e.g., "task")
field	str	TypeDB attribute name (e.g., "display_id")
variant	str	DTO variant to override ("create", "out", "patch")
required	bool | None	Override requiredness. None keeps the schema default.
default	str | None	Override default value as a Python literal string. None keeps the schema default.

FieldOverride dataclass

FieldOverride(required=None, default=None)

Per-variant override for a single field's requiredness or default.

Use this inside BaseClassConfig.create_field_overrides to make a normally-required field optional (or vice-versa) on a specific DTO variant.

Attributes:

Name	Type	Description
required	bool | None	Override requiredness. None keeps the schema default.
default	str | None	Override default value as a Python literal string (e.g., "None"). None keeps the schema default.

FieldSyncConfig dataclass

FieldSyncConfig(field_a, field_b)

Configuration for syncing two fields in a model validator.

When one field is set but not the other, the value is copied.

Attributes:

Name	Type	Description
field_a	str	First field name
field_b	str	Second field name

Example

FieldSyncConfig(field_a="description", field_b="content")

Generates a model_validator that syncs description <-> content

ValidatorConfig dataclass

ValidatorConfig(name, pattern=None)

Configuration for a custom Pydantic validator type.

Validators are rendered as Annotated types with AfterValidator.

Attributes:

Name	Type	Description
name	str	The validator type name (e.g., "DisplayId")
pattern	str | None	Regex pattern for validation (optional). For complex validators that need more than a regex pattern, use preamble in DTOConfig to define custom functions.

Example

ValidatorConfig(name="DisplayId", pattern=r"^[A-Z]{1,10}-\d+$")

Generates:

def _validate_display_id(value: str) -> str:

if not re.match(r"^[A-Z]{1,10}-\d+$", value):

raise ValueError(...)

return value

DisplayId = Annotated[str, AfterValidator(_validate_display_id)]

ParsedSchema dataclass

ParsedSchema(attributes=dict(), entities=dict(), relations=dict(), functions=dict(), structs=dict())

Container for all parsed schema components.

This is the main output of the parser and input to the renderers.

accumulate_inheritance

accumulate_inheritance()

Propagate inherited members down all type hierarchies.

Source code in type_bridge/generator/models.py

def accumulate_inheritance(self) -> None:
    """Propagate inherited members down all type hierarchies."""
    _accumulate_entity_inheritance(self)
    _accumulate_relation_inheritance(self)

parse_tql_schema

parse_tql_schema(schema_content)

Parse a TQL schema string into a :class:ParsedSchema.

Uses the Rust TypeSchema parser when available for speed, falling back to the Lark-based parser otherwise. Annotations and docstrings (extracted from comments) are always applied on the Python side.

Source code in type_bridge/generator/parser.py

def parse_tql_schema(schema_content: str) -> ParsedSchema:
    """Parse a TQL schema string into a :class:`ParsedSchema`.

    Uses the Rust ``TypeSchema`` parser when available for speed, falling
    back to the Lark-based parser otherwise.  Annotations and docstrings
    (extracted from comments) are always applied on the Python side.
    """
    # Step 1: Always extract annotations from comments first
    entity_annots, attr_annots, rel_annots, role_annots = extract_annotations(schema_content)

    # Step 2: Try Rust parser
    if _RUST_SCHEMA_AVAILABLE:
        try:
            rust_schema = _RustTypeSchema.from_typeql(schema_content)  # type: ignore[union-attr]
            schema = _rust_schema_to_parsed(
                rust_schema, entity_annots, attr_annots, rel_annots, role_annots
            )
            logger.debug("Parsed schema using Rust core")

            # Step 3: If functions/structs present, supplement with Lark
            if _HAS_FUN_OR_STRUCT.search(schema_content):
                logger.debug("Schema contains functions/structs, supplementing with Lark parser")
                lark_schema = _parse_with_lark(
                    schema_content, entity_annots, attr_annots, rel_annots, role_annots
                )
                schema.functions = lark_schema.functions
                schema.structs = lark_schema.structs

            return schema
        except Exception:
            logger.warning("Rust parser failed, falling back to Lark parser", exc_info=True)

    # Step 4: Full Lark fallback
    if not _RUST_SCHEMA_AVAILABLE:
        logger.debug("Rust core not available, using Lark parser")
    return _parse_with_lark(schema_content, entity_annots, attr_annots, rel_annots, role_annots)

generate_models

generate_models(schema, output_dir, *, implicit_key_attributes=None, schema_version='1.0.0', copy_schema=True, schema_path=None, generate_dto=False, dto_config=None)

Generate TypeBridge models from a TypeDB schema.

Parameters:

Name	Type	Description	Default
schema	str | Path	Either a path to a .tql file, or the schema text directly	required
output_dir	str | Path	Directory to write the generated package to	required
implicit_key_attributes	Iterable[str] | None	Attribute names to treat as @key even if not marked	None
schema_version	str	Version string for SCHEMA_VERSION constant	'1.0.0'
copy_schema	bool	Whether to copy the schema file to the output directory	True
schema_path	str | Path | None	Custom path for the schema file. If relative, resolved against output_dir. If None and copy_schema=True, uses "schema.tql" in output_dir.	None
generate_dto	bool	Whether to generate Pydantic API DTOs	False
dto_config	DTOConfig | None	Configuration for DTO generation (custom base classes, validators, etc.)	None

Source code in type_bridge/generator/__init__.py

def generate_models(
    schema: str | Path,
    output_dir: str | Path,
    *,
    implicit_key_attributes: Iterable[str] | None = None,
    schema_version: str = "1.0.0",
    copy_schema: bool = True,
    schema_path: str | Path | None = None,
    generate_dto: bool = False,
    dto_config: DTOConfig | None = None,
) -> None:
    """Generate TypeBridge models from a TypeDB schema.

    Args:
        schema: Either a path to a .tql file, or the schema text directly
        output_dir: Directory to write the generated package to
        implicit_key_attributes: Attribute names to treat as @key even if not marked
        schema_version: Version string for SCHEMA_VERSION constant
        copy_schema: Whether to copy the schema file to the output directory
        schema_path: Custom path for the schema file. If relative, resolved against
            output_dir. If None and copy_schema=True, uses "schema.tql" in output_dir.
        generate_dto: Whether to generate Pydantic API DTOs
        dto_config: Configuration for DTO generation (custom base classes, validators, etc.)
    """
    # Resolve schema text
    schema_source_path: Path | None = None
    if isinstance(schema, Path):
        schema_source_path = schema
    elif isinstance(schema, str):
        # Check if it looks like a file path (short string, no newlines)
        if len(schema) < 500 and "\n" not in schema:
            try:
                candidate = Path(schema)
                if candidate.exists() and candidate.is_file():
                    schema_source_path = candidate
            except OSError:
                pass  # Not a valid path

    if schema_source_path:
        schema_text = schema_source_path.read_text(encoding="utf-8")
    else:
        schema_text = str(schema)

    # Create output directory
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)

    # Parse schema
    parsed = parse_tql_schema(schema_text)
    implicit_keys = set(implicit_key_attributes or set())

    # Build class name mappings
    attr_class_names = build_class_name_map(parsed.attributes)
    entity_class_names = build_class_name_map(parsed.entities)
    relation_class_names = build_class_name_map(parsed.relations)
    struct_class_names = build_class_name_map(parsed.structs)

    # Generate and write files
    (output / "attributes.py").write_text(
        render_attributes(parsed, attr_class_names),
        encoding="utf-8",
    )

    (output / "entities.py").write_text(
        render_entities(parsed, attr_class_names, entity_class_names, implicit_keys),
        encoding="utf-8",
    )

    (output / "relations.py").write_text(
        render_relations(parsed, attr_class_names, entity_class_names, relation_class_names),
        encoding="utf-8",
    )

    # Render functions if present
    functions_content = render_functions(parsed)
    functions_present = False
    if functions_content:
        (output / "functions.py").write_text(functions_content, encoding="utf-8")
        functions_present = True

    # Render structs if present
    structs_content = render_structs(parsed, struct_class_names)
    if structs_content:
        (output / "structs.py").write_text(structs_content, encoding="utf-8")

    # Render registry with pre-computed metadata
    (output / "registry.py").write_text(
        render_registry(
            parsed,
            attr_class_names,
            entity_class_names,
            relation_class_names,
            schema_version=schema_version,
            schema_text=schema_text,
        ),
        encoding="utf-8",
    )

    # Render Pydantic DTOs if requested
    if generate_dto:
        (output / "api_dto.py").write_text(
            render_api_dto(parsed, config=dto_config),
            encoding="utf-8",
        )

    # Determine schema output location
    schema_filename: str | None = None
    schema_output_path: Path | None = None
    if copy_schema:
        if schema_path is None:
            # Default: schema.tql in output directory
            schema_output_path = output / "schema.tql"
            schema_filename = "schema.tql"
        else:
            resolved_path = Path(schema_path)
            if resolved_path.is_absolute():
                schema_output_path = resolved_path
                # Only include loader if schema is in the output directory
                try:
                    resolved_path.relative_to(output.resolve())
                    schema_filename = resolved_path.name
                except ValueError:
                    schema_filename = None  # Outside output dir, no loader
            else:
                # Relative path - resolve against output dir
                schema_output_path = output / resolved_path
                # Only include loader if it's a simple filename (no subdirs)
                if resolved_path.parent == Path("."):
                    schema_filename = str(resolved_path)
                else:
                    schema_filename = None  # In subdir, loader won't work

    (output / "__init__.py").write_text(
        render_package_init(
            attr_class_names,
            entity_class_names,
            relation_class_names,
            schema_version=schema_version,
            include_schema_loader=schema_filename is not None,
            schema_filename=schema_filename,
            functions_present=functions_present,
        ),
        encoding="utf-8",
    )

    # Copy schema file if requested
    if schema_output_path:
        schema_output_path.parent.mkdir(parents=True, exist_ok=True)
        schema_output_path.write_text(schema_text, encoding="utf-8")