Code Generation

The kumoai.codegen module can generate Python SDK code for supported connectors, tables, graphs, and predictive queries. Use it to turn a supported saved entity into a reproducible script that can be reviewed, versioned, and rerun.

Generate code by ID

Initialize the SDK, then call generate_code() with an entity ID. When the entity class is ambiguous, include entity_class.

1
import os
2
import kumoai as kumo
3
from kumoai.codegen import generate_code
4
5
kumo.init(
6
    url=os.environ["KUMO_API_ENDPOINT"],
7
    api_key=os.environ["KUMO_API_KEY"],
8
)
9
10
code = generate_code({
11
    "id": "myconnector",
12
    "entity_class": "S3Connector",
13
})
14
15
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.

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

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

Generate code from an in-memory object

For advanced workflows, pass an SDK object directly with the object input mode.

1
code = generate_code({"object": graph})
2
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

The SDK includes a CLI module. Invoke it 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 your environment exposes the optional kumo-codegen entry point, the equivalent command is:

$
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.

1
from kumoai.codegen import (
2
    CyclicDependencyError,
3
    UnsupportedEntityError,
4
    generate_code,
5
)
6
7
try:
8
    code = generate_code({"id": entity_id, "entity_class": entity_class})
9
except UnsupportedEntityError as exc:
10
    print("This entity type is not supported by codegen:", exc)
11
except CyclicDependencyError as exc:
12
    print("The entity graph contains a cycle:", exc)
13
except ValueError as exc:
14
    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 UI or notebook changes, then review the diff before replacing a production script.