kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/codegen/generate.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import logging
+import os
+from collections import OrderedDict
+from typing import Any, Optional
+
+from kumoai.codegen.context import CodegenContext
+from kumoai.codegen.exceptions import (
+    CyclicDependencyError,
+    UnsupportedEntityError,
+)
+from kumoai.codegen.identity import get_config_id
+from kumoai.codegen.loader import load_from_id
+from kumoai.codegen.naming import NameManager
+from kumoai.codegen.registry import (
+    REG,
+    Handler,
+    execute_in_env,
+    get_from_env,
+    init_execution_env,
+    register_shared_parents,
+    store_object_var,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _get_handler(obj_type: type) -> Handler:
+    if obj_type not in REG:
+        raise UnsupportedEntityError(
+            f"No handler registered for type: {obj_type.__name__}")
+    return REG[obj_type]
+
+
+def get_kumo_id(obj: object) -> str:
+    if hasattr(obj, 'id'):
+        return obj.id
+    elif hasattr(obj, 'name'):
+        return obj.name
+    elif hasattr(obj, 'source_name'):
+        return obj.source_name
+    else:
+        return ''
+
+
+def _generate(
+    obj: object,
+    name_manager: NameManager,
+    config_to_var: dict[str, str],
+    stack: set[int],
+    codegen_ctx: CodegenContext,
+    context: Optional[dict[str, Any]] = None,
+    id_to_var: Optional[dict[int, str]] = None,
+) -> tuple[list[str], list[str]]:
+    """Generate code for an object and its parents.
+
+    Args:
+        obj: The object to generate code for.
+        name_manager: The name manager to use for variable names.
+        config_to_var: A dictionary mapping config IDs to variable names.
+        stack: A set of object IDs to detect cycles.
+        codegen_ctx: The codegen context.
+        context: A dictionary of context information.
+        id_to_var: A dictionary mapping object IDs to variable names.
+
+    Returns:
+        A tuple containing a list of imports and a list of lines of code.
+    """
+    if id_to_var is None:
+        id_to_var = {}
+
+    obj_id = id(obj)
+    config_id = get_config_id(obj)
+
+    # Check for configuration-based deduplication first
+    if config_id in config_to_var:
+        # Reuse existing variable for this configuration
+        id_to_var[obj_id] = config_to_var[config_id]
+        return [], []
+
+    # Check for exact object reuse (faster path)
+    if obj_id in id_to_var:
+        return [], []
+
+    # Cycle detection using real object IDs
+    if obj_id in stack:
+        raise CyclicDependencyError(
+            f"Cyclic dependency detected for object ID: {obj_id}")
+
+    stack.add(obj_id)
+    handler = _get_handler(type(obj))
+    all_imports, all_lines = [], []
+
+    for parent in handler.parents(obj, codegen_ctx):
+        parent_imports, parent_lines = _generate(parent, name_manager,
+                                                 config_to_var, stack,
+                                                 codegen_ctx, context,
+                                                 id_to_var)
+        all_imports.extend(parent_imports)
+        all_lines.extend(parent_lines)
+
+    # Register shared parents if handler supports it
+    register_shared_parents(codegen_ctx, obj, handler)
+
+    var_name = name_manager.assign_entity_variable(obj)
+    # Store both config-based and id-based mappings
+    config_to_var[config_id] = var_name
+    id_to_var[obj_id] = var_name
+
+    # Store in codegen context for handlers to access
+    store_object_var(codegen_ctx, obj, var_name)
+
+    all_imports.extend(handler.required_imports(obj))
+    context = context or {}
+    context['target_id'] = get_kumo_id(obj)
+    creation_lines = handler.emit_lines(obj, var_name, context, codegen_ctx)
+
+    # Execute creation lines immediately in context environment
+    execute_in_env(codegen_ctx, creation_lines, handler.required_imports(obj))
+
+    all_lines.extend(creation_lines)
+
+    if handler.detect_edits:
+        # Get baseline object from context environment
+        baseline_obj = get_from_env(codegen_ctx, var_name)
+        if baseline_obj is not None:
+            edits = handler.detect_edits(obj, baseline_obj, name_manager)
+            for edit in edits:
+                edit_lines = edit.emit_lines(var_name)
+                all_lines.extend(edit_lines)
+                all_imports.extend(edit.required_imports)
+                # Execute edit lines immediately in context environment
+                execute_in_env(codegen_ctx, edit_lines, edit.required_imports)
+
+    stack.remove(obj_id)
+    return all_imports, all_lines
+
+
+def _load_entity_from_spec(input_spec: dict[str, Any]) -> object:
+    """Load entity from input specification."""
+    if "id" in input_spec:
+        entity_class = input_spec.get("entity_class")
+        return load_from_id(input_spec["id"], entity_class)
+    elif "json" in input_spec:
+        raise NotImplementedError("JSON loading not yet implemented")
+    elif "object" in input_spec:
+        return input_spec["object"]
+    else:
+        raise ValueError("input_spec must contain 'id', 'json', or 'object'")
+
+
+def _write_script(code: str, output_path: str) -> None:
+    """Write generated code to file."""
+    with open(output_path, "w") as f:
+        f.write(code)
+
+
+def _assemble_code(imports: list[str], lines: list[str]) -> str:
+    """Assemble final code from components."""
+    from kumoai import __version__
+
+    header = [
+        f"# Generated with Kumo SDK version: {__version__}",
+        "import kumoai as kumo",
+        "import os",
+        "",
+        'kumo.init(url=os.getenv("KUMO_API_ENDPOINT"), '
+        'api_key=os.getenv("KUMO_API_KEY"))',
+        "",
+    ]
+
+    unique_imports = list(OrderedDict.fromkeys(imports))
+    code = header + unique_imports + [""] + lines
+    return "\n".join(code) + "\n"
+
+
+def _init_kumo() -> None:
+    """Initialize Kumo SDK for this python session."""
+    import kumoai as kumo
+    if os.getenv("KUMO_API_ENDPOINT") is None:
+        logger.warning("KUMO_API_ENDPOINT env variable is not set, "
+                       "assuming kumo.init has already been called")
+        return
+    if os.getenv("KUMO_API_KEY") is None:
+        logger.warning("KUMO_API_KEY env variable is not set, "
+                       "assuming kumo.init has already been called")
+        return
+    kumo.init(url=os.getenv("KUMO_API_ENDPOINT"),
+              api_key=os.getenv("KUMO_API_KEY"))
+
+
+def generate_code(input_spec: dict[str, Any],
+                  output_path: Optional[str] = None) -> str:
+    """Generate Python SDK code from Kumo entity specification."""
+    # Create codegen context for this generation session
+    codegen_ctx = CodegenContext()
+
+    # Initialize execution environment in context
+    init_execution_env(codegen_ctx)
+
+    _init_kumo()
+    entity = _load_entity_from_spec(input_spec)
+
+    context = {}
+    if "id" in input_spec:
+        context["input_method"] = "id"
+        context["target_id"] = input_spec["id"]
+    elif "json" in input_spec:
+        context["input_method"] = "json"
+    else:
+        context["input_method"] = "object"
+
+    name_manager = NameManager()
+    imports, lines = _generate(entity, name_manager, config_to_var={},
+                               stack=set(), codegen_ctx=codegen_ctx,
+                               context=context, id_to_var={})
+
+    code = _assemble_code(imports, lines)
+    if output_path:
+        _write_script(code, output_path)
+    return code

kumoai/codegen/handlers/__init__.py

@@ -0,0 +1,4 @@
+from . import connector  # noqa
+from . import table  # noqa
+from . import graph  # noqa
+from . import pquery  # noqa

kumoai/codegen/handlers/connector.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from typing import Dict, Type
+
+from kumoai.codegen.context import CodegenContext
+from kumoai.codegen.handlers.utils import _get_canonical_import_path
+from kumoai.codegen.registry import Handler
+from kumoai.connector import (
+    BigQueryConnector,
+    DatabricksConnector,
+    FileUploadConnector,
+    S3Connector,
+    SnowflakeConnector,
+)
+
+
+def _get_by_name_handler_factory(cls: type) -> Handler:
+    """Factory for creating a Handler for any connector that uses the
+    get_by_name pattern.
+    """
+    def get_imports(obj: object, bound_cls: type = cls) -> list[str]:
+        canonical_module = _get_canonical_import_path(bound_cls)
+        return [f"from {canonical_module} import {bound_cls.__name__}"]
+
+    def get_lines(
+        obj: object,
+        var_name: str,
+        context: dict,
+        codegen_ctx: CodegenContext,
+    ) -> list[str]:
+        assert isinstance(obj, cls)
+        assert hasattr(obj, "name")
+        obj_name = getattr(obj, "name")
+        return [f"{var_name} = {cls.__name__}.get_by_name('{obj_name}')"]
+
+    return Handler(
+        parents=lambda e, ctx: [],
+        required_imports=get_imports,
+        emit_lines=get_lines,
+        detect_edits=None,
+    )
+
+
+def get_handlers() -> Dict[Type, Handler]:
+    """Returns a dictionary of handlers for all connector types."""
+    handlers: Dict[Type, Handler] = {}
+
+    # S3Connector gets special handling to support both patterns
+
+    def _s3_connector_parents(obj: object,
+                              codegen_ctx: CodegenContext) -> list[object]:
+        return []
+
+    def _s3_connector_imports(obj: object) -> list[str]:
+        canonical_module = _get_canonical_import_path(S3Connector)
+        return [f"from {canonical_module} import S3Connector"]
+
+    def _s3_connector_emit_lines(
+        obj: object,
+        var_name: str,
+        context: dict,
+        codegen_ctx: CodegenContext,
+    ) -> list[str]:
+        assert isinstance(obj, S3Connector)
+        assert hasattr(obj, "name")
+
+        # Check if connector has root_dir attribute and it's not None
+        if hasattr(obj, "root_dir") and obj.root_dir is not None:
+            root_dir = getattr(obj, "root_dir")
+            return [f"{var_name} = S3Connector(root_dir='{root_dir}')"]
+        else:
+            obj_name = getattr(obj, "name")
+            return [f"{var_name} = S3Connector.get_by_name('{obj_name}')"]
+
+    handlers[S3Connector] = Handler(
+        parents=_s3_connector_parents,
+        required_imports=_s3_connector_imports,
+        emit_lines=_s3_connector_emit_lines,
+        detect_edits=None,
+    )
+
+    # Other persistent connectors use the standard get_by_name pattern
+    other_persistent_connectors = [
+        BigQueryConnector,
+        SnowflakeConnector,
+        DatabricksConnector,
+    ]
+    for connector_cls in other_persistent_connectors:
+        handlers[connector_cls] = _get_by_name_handler_factory(connector_cls)
+
+    def _file_upload_parents(obj: object,
+                             codegen_ctx: CodegenContext) -> list[object]:
+        return []
+
+    def _file_upload_imports(obj: object) -> list[str]:
+        canonical_module = _get_canonical_import_path(type(obj))
+        return [f"from {canonical_module} import {type(obj).__name__}"]
+
+    def _file_upload_emit_lines(
+        obj: object,
+        var_name: str,
+        context: dict,
+        codegen_ctx: CodegenContext,
+    ) -> list[str]:
+        assert isinstance(obj, FileUploadConnector)
+        return [
+            f"{var_name} = {type(obj).__name__}"
+            f"(file_type='{obj.file_type}')"
+        ]
+
+    handlers[FileUploadConnector] = Handler(
+        parents=_file_upload_parents,
+        required_imports=_file_upload_imports,
+        emit_lines=_file_upload_emit_lines,
+        detect_edits=None,
+    )
+
+    return handlers

kumoai/codegen/handlers/graph.py

@@ -0,0 +1,71 @@
+import logging
+from typing import Dict, Type
+
+from kumoai.codegen.context import CodegenContext
+from kumoai.codegen.handlers.utils import _get_canonical_import_path
+from kumoai.codegen.registry import Handler, get_object_var
+from kumoai.graph import Edge, Graph
+
+logger = logging.getLogger(__name__)
+
+
+def get_handlers() -> Dict[Type, Handler]:
+    """Returns a dictionary of handlers for Graph types."""
+    handlers: Dict[Type, Handler] = {}
+
+    def _graph_parents(obj: object,
+                       codegen_ctx: CodegenContext) -> list[object]:
+        """Graph depends on its tables and edges."""
+        assert isinstance(obj, Graph)
+        return [table for table in obj.tables.values()]
+
+    def _graph_imports(obj: object) -> list[str]:
+        """Get import statements needed for Graph."""
+        imports_needed = [Graph, Edge]
+        imports = []
+        for obj_type in imports_needed:
+            canonical_module = _get_canonical_import_path(obj_type)
+            imports.append(
+                f"from {canonical_module} import {obj_type.__name__}")
+        return imports
+
+    def _graph_emit_lines(
+        obj: object,
+        var_name: str,
+        context: dict,
+        codegen_ctx: CodegenContext,
+    ) -> list[str]:
+        """Generate code lines to recreate a Graph using Graph()."""
+        assert isinstance(obj, Graph)
+
+        tables_vars = {
+            table_name: get_object_var(codegen_ctx, table)
+            for table_name, table in obj.tables.items()
+        }
+        tables_format_inner = ", ".join(f"{key!r}: {value}"
+                                        for key, value in tables_vars.items())
+        all_tables_var_name = f"{var_name}_tables"
+        all_edges_var_name = f"{var_name}_edges"
+
+        note = (f"# Note: This could also be done with Graph.load("
+                f"'{context['target_id']}') for simpler code")
+
+        lines = [
+            note, f"{all_tables_var_name} = {{{tables_format_inner}}}",
+            f"{all_edges_var_name} = {obj.edges}",
+            f"{var_name} = Graph(tables={all_tables_var_name}, "
+            f"edges={all_edges_var_name})", f"{var_name}.validate()",
+            f"# Optionally, you can save the graph to backend using "
+            f"{var_name}.save({context['target_id']}) and then load it"
+        ]
+
+        return lines
+
+    handlers[Graph] = Handler(
+        parents=_graph_parents,
+        required_imports=_graph_imports,
+        emit_lines=_graph_emit_lines,
+        detect_edits=None,
+    )
+
+    return handlers

kumoai/codegen/handlers/pquery.py

@@ -0,0 +1,62 @@
+import logging
+from typing import Dict, Type
+
+from kumoai.codegen.context import CodegenContext
+from kumoai.codegen.handlers.utils import _get_canonical_import_path
+from kumoai.codegen.registry import Handler, get_object_var
+from kumoai.pquery import PredictiveQuery
+
+logger = logging.getLogger(__name__)
+
+
+def get_handlers() -> Dict[Type, Handler]:
+    """Returns a dictionary of handlers for PredictiveQuery types."""
+    handlers: Dict[Type, Handler] = {}
+
+    def _pquery_parents(obj: object,
+                        codegen_ctx: CodegenContext) -> list[object]:
+        """PredictiveQuery depends on its graph."""
+        assert isinstance(obj, PredictiveQuery)
+        return [obj.graph]
+
+    def _pquery_imports(obj: object) -> list[str]:
+        """Get import statements needed for PredictiveQuery."""
+        imports_needed = [PredictiveQuery]
+        imports = []
+        for obj_type in imports_needed:
+            canonical_module = _get_canonical_import_path(obj_type)
+            imports.append(
+                f"from {canonical_module} import {obj_type.__name__}")
+        return imports
+
+    def _pquery_emit_lines(
+        obj: object,
+        var_name: str,
+        context: dict,
+        codegen_ctx: CodegenContext,
+    ) -> list[str]:
+        """Generate code lines to recreate a PredictiveQuery."""
+        assert isinstance(obj, PredictiveQuery)
+
+        graph_var = get_object_var(codegen_ctx, obj.graph)
+
+        if '\n' in obj.query:
+            formatted_query = f'"""{obj.query}"""'
+        else:
+            formatted_query = repr(obj.query)
+
+        lines = [
+            f"{var_name} = PredictiveQuery(" + f"graph={graph_var}," +
+            f" query={formatted_query}," + ")", f"{var_name}.validate()"
+        ]
+
+        return lines
+
+    handlers[PredictiveQuery] = Handler(
+        parents=_pquery_parents,
+        required_imports=_pquery_imports,
+        emit_lines=_pquery_emit_lines,
+        detect_edits=None,
+    )
+
+    return handlers

kumoai/codegen/handlers/table.py

@@ -0,0 +1,109 @@
+import logging
+from typing import Dict, Sequence, Type
+
+from kumoai.codegen.context import CodegenContext
+from kumoai.codegen.edits import (
+    UniversalReplacementEdit,
+    detect_edits_recursive,
+)
+from kumoai.codegen.handlers.utils import _get_canonical_import_path
+from kumoai.codegen.naming import NameManager
+from kumoai.codegen.registry import Handler, get_object_var
+from kumoai.graph.table import Table
+
+logger = logging.getLogger(__name__)
+
+
+def get_handlers() -> Dict[Type, Handler]:
+    """Returns a dictionary of handlers for Table types."""
+    handlers: Dict[Type, Handler] = {}
+
+    def _table_parents(obj: object,
+                       codegen_ctx: CodegenContext) -> list[object]:
+        """Table depends on its source_table's connector."""
+        assert isinstance(obj, Table)
+        return [obj.source_table.connector]
+
+    def _table_imports(obj: object) -> list[str]:
+        """Get import statements needed for Table."""
+        canonical_module = _get_canonical_import_path(Table)
+        return [f"from {canonical_module} import Table"]
+
+    def _build_args(obj: Table, connector_var: str,
+                    table_name: str) -> list[str]:
+        """Build arguments for Table.from_source_table() call."""
+        args = [f"source_table={connector_var}['{table_name}']"]
+
+        if obj.primary_key is not None:
+            args.append(f"primary_key='{obj.primary_key.name}'")
+
+        if obj.time_column is not None:
+            args.append(f"time_column='{obj.time_column.name}'")
+
+        if obj.end_time_column is not None:
+            args.append(f"end_time_column='{obj.end_time_column.name}'")
+
+        source_table_cols = {col.name: col for col in obj.source_table.columns}
+        table_cols = {col.name: col for col in obj.columns}
+
+        if len(source_table_cols) != len(table_cols):
+            assert len(source_table_cols) > len(table_cols)
+            table_col_names = [col_name for col_name in table_cols.keys()]
+            args.append(f"column_names={table_col_names}")
+
+        return args
+
+    def _table_emit_lines(
+        obj: object,
+        var_name: str,
+        context: dict,
+        codegen_ctx: CodegenContext,
+    ) -> list[str]:
+        """Generate code lines to recreate a Table using from_source_table."""
+        assert isinstance(obj, Table)
+
+        connector_var = get_object_var(codegen_ctx, obj.source_table.connector)
+        table_name = obj.source_table.name
+
+        args = _build_args(obj, connector_var, table_name)
+
+        lines = []
+
+        if context.get("input_method") == "id" and context.get("target_id"):
+            note = (f"# Note: This could also be done with Table.load("
+                    f"'{context['target_id']}') for simpler code")
+            lines.append(note)
+
+        if len(args) == 1:
+            lines.append(f"{var_name} = Table.from_source_table({args[0]})")
+        else:
+            args_str = ",\n    ".join(args)
+            lines.append(
+                f"{var_name} = Table.from_source_table(\n    {args_str},\n)")
+
+        return lines
+
+    def _table_detect_edits(
+            target: object, baseline: object,
+            name_manager: NameManager) -> Sequence[UniversalReplacementEdit]:
+        """Detect edits needed to make baseline match target table."""
+        assert isinstance(target, Table)
+        assert isinstance(baseline, Table)
+
+        try:
+            result = detect_edits_recursive(target, baseline, "", name_manager)
+            logger.debug(f"Found for table {len(result.edits)} edits with "
+                         f"{len(result.imports)} imports")
+            return result.edits
+        except Exception as e:
+            logger.error(f"Error during table edit detection: {e}")
+            return []
+
+    handlers[Table] = Handler(
+        parents=_table_parents,
+        required_imports=_table_imports,
+        emit_lines=_table_emit_lines,
+        detect_edits=_table_detect_edits,
+    )
+
+    return handlers

kumoai/codegen/handlers/utils.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import importlib
+from typing import Type
+
+
+def _get_canonical_import_path(cls: Type) -> str:
+    """Dynamically finds the shortest, most canonical import path.
+
+    For example, given the S3Connector class, it would prefer
+    'from kumoai import S3Connector' over
+    'from kumoai.connector import S3Connector' over
+    'from kumoai.connector.s3_connector import S3Connector'.
+    """
+    base_module_path = cls.__module__
+    class_name = cls.__name__
+
+    parts = base_module_path.split(".")
+
+    # The longest path is the one where the class is defined
+    canonical_path = base_module_path
+
+    # Walk up the module hierarchy to find a shorter path.
+    # e.g., from 'kumoai.connector.s3_connector' to 'kumoai.connector'
+    for i in range(len(parts) - 1, 0, -1):
+        parent_module_path = ".".join(parts[:i])
+        try:
+            parent_module = importlib.import_module(parent_module_path)
+            if hasattr(parent_module, class_name):
+                if getattr(parent_module, class_name) is cls:
+                    canonical_path = parent_module_path
+                else:
+                    # A different object has the same name
+                    break
+            else:
+                # The class isn't in this parent module
+                break
+        except ImportError:
+            # If we can't import a parent, we can't go higher.
+            break
+
+    return canonical_path