Code Generation

The kumoai.codegen module generates Python SDK code for supported connectors, tables, graphs, and predictive queries. Use it to export any supported Kumo entity as a reproducible Python script you can review, version, and rerun.

Generate code by ID

Initialize the SDK, then call generate_code() with an entity ID. When the entity ID alone does not identify the type, also pass entity_class.

import os
import kumoai as kumo
from kumoai.codegen import generate_code
kumo.init(
    url=os.environ["KUMO_API_ENDPOINT"],
    api_key=os.environ["KUMO_API_KEY"],
)
code = generate_code({
    "id": "myconnector",
    "entity_class": "S3Connector",
})
print(code)

The generated script includes Kumo initialization boilerplate and the imports required to recreate the entity and its dependencies.

Write generated code to a file

Pass output_path to write the generated script directly.

generate_code(
    {"id": "my-predictive-query", "entity_class": "PredictiveQuery"},
    output_path="recreate_predictive_query.py",
)

Review the generated file before running it in production, especially if it will create or update Kumo entities in a shared workspace.

Generate code from an in-memory object

Pass an in-memory SDK object directly using the object input mode.

code = generate_code({"object": graph})
print(code)

This is useful in notebooks when you have modified an object interactively and want to export the equivalent Python code.

Command-line usage

Invoke the CLI module directly with python -m:

python -m kumoai.codegen.cli --id myconnector --entity-class S3Connector -o connector.py
python -m kumoai.codegen.cli --id my-predictive-query --entity-class PredictiveQuery

If the optional kumo-codegen entry point is available in your environment:

kumo-codegen --id myconnector --entity-class S3Connector --output connector.py

Use --verbose to print additional diagnostics.

Supported input modes

How dependency generation works

When you generate code for an entity with dependencies, codegen emits parent objects first. For example, generating a predictive query can emit the connector, source tables, Kumo tables, and graph needed by that query.

The generator deduplicates equivalent supported connector configurations; other entities are reused only when the same object instance is encountered.

Error handling

Catch codegen-specific exceptions when building automation around generated scripts.

from kumoai.codegen import (
    CyclicDependencyError,
    UnsupportedEntityError,
    generate_code,
)
try:
    code = generate_code({"id": entity_id, "entity_class": entity_class})
except UnsupportedEntityError as exc:
    print("This entity type is not supported by codegen:", exc)
except CyclicDependencyError as exc:
    print("The entity graph contains a cycle:", exc)
except ValueError as exc:
    print("The ID or entity class is invalid or unsupported:", exc)

Best practices

• Commit generated scripts alongside the analysis or pipeline that depends on them.
• Treat generated code as a starting point; rename variables and add comments before production use.
• Prefer environment variables for KUMO_API_ENDPOINT and KUMO_API_KEY so generated scripts do not contain secrets.
• Regenerate code after substantial changes in the Kumo UI or notebooks, then review the diff before replacing a production script.