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

kumoai/codegen/identity.py

@@ -0,0 +1,114 @@
+"""Configuration-based identity system for codegen deduplication."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+from kumoai.connector import (
+    BigQueryConnector,
+    DatabricksConnector,
+    FileUploadConnector,
+    S3Connector,
+    SnowflakeConnector,
+)
+
+
+def get_config_id(obj: Any) -> str:
+    """Return configuration-based identity for codegen deduplication.
+
+    Objects with identical configurations get the same config_id, enabling
+    variable reuse during code generation. Uses get_editable_attributes() to
+    automatically capture all configurable properties.
+
+    This configuration-based identity is used for deduplication only of
+    object which are SAFE to be shared
+    when they have the same configuration always.
+    For example, a S3Connector
+    with the same root_dir and name can always be shared,
+    it does not matter if some other objects wants to copy it.
+
+    But for certain objects like Graphs,
+    some object may want 2 copies of the same graph,
+    and there we should not use this config_id
+    and always use the memory address.
+    For example, a Graph with the same tables and edges can always be shared,
+    it does not matter if some other objects wants to copy it.
+
+    Args:
+        obj: Object to get configuration ID for
+
+    Returns:
+        Configuration-based identity string
+
+    Note:
+        - For deduplication only, not cycle detection (use id() for cycles)
+        - Only applied to connector types for now; other objects use memory ID
+    """
+    generic_object_types = (S3Connector, BigQueryConnector, SnowflakeConnector,
+                            DatabricksConnector, FileUploadConnector)
+    if isinstance(obj, generic_object_types):
+        return _get_generic_config_id(obj)
+    else:
+        return f"id_{id(obj)}"
+
+
+def _get_generic_config_id(obj: Any) -> str:
+    """Generate config ID by hashing object type and editable attributes.
+
+    Uses get_editable_attributes() to capture all configurable properties,
+    then creates a SHA256 hash for consistent identity.
+    """
+    try:
+        # Import here to avoid circular imports
+        from kumoai.codegen.edits import get_editable_attributes
+
+        # Get object type name
+        obj_type = type(obj).__name__
+
+        # Get all editable attributes
+        editable_attrs = get_editable_attributes(obj)
+
+        # Build configuration dict
+        config: dict[str, Any] = {'type': obj_type, 'attributes': {}}
+
+        # Extract values for all editable attributes
+        for attr_name in sorted(editable_attrs):  # Sort for consistent hashing
+            try:
+                attr_value = getattr(obj, attr_name)
+                config['attributes'][attr_name] = _serialize_value(attr_value)
+            except (AttributeError, RuntimeError, TypeError):
+                # Skip attributes that can't be accessed
+                continue
+
+        # Create hash from configuration
+        config_str = json.dumps(config, sort_keys=True)
+        config_hash = hashlib.sha256(config_str.encode()).hexdigest()[:16]
+
+        return f"{obj_type}_{config_hash}"
+
+    except Exception:
+        # Fallback to memory address if hashing fails
+        return f"{type(obj).__name__}_{id(obj)}"
+
+
+def _serialize_value(value: Any) -> Any:
+    """Convert value to JSON-serializable format for consistent hashing.
+
+    Handles nested objects by recursively applying config-based identity.
+    """
+    if value is None:
+        return None
+    elif isinstance(value, (str, int, float, bool)):
+        return value
+    elif isinstance(value, (list, tuple)):
+        return [_serialize_value(item) for item in value]
+    elif isinstance(value, dict):
+        return {k: _serialize_value(v) for k, v in value.items()}
+    elif hasattr(value, '__dict__'):
+        # For objects with __dict__, recurse into their config_id
+        return get_config_id(value)
+    else:
+        # For other types, convert to string
+        return str(value)

kumoai/codegen/loader.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import kumoai as kumo
+from kumoai.exceptions import HTTPException
+
+
+def _get_supported_entities() -> dict[str, type]:
+    """Get all supported entity types from registered handlers."""
+    from kumoai.codegen.registry import REG
+
+    return {cls.__name__: cls for cls in REG.keys()}
+
+
+# This map is for converting string names to Python types.
+# The keys should be the string a user would provide.
+# Auto-generated from handlers registry while serving as supported
+# entities filter.
+ENTITY_TYPE_MAP = _get_supported_entities()
+
+# This map contains prefixes for type inference from their ID.
+ENTITY_PREFIX_MAP = {
+    # Job Types
+    "gen-traintable-job": kumo.TrainingTable,
+    "trainingjob": kumo.TrainingJob,
+    "bp-job": kumo.BatchPredictionJob,
+    # Query Types
+    "pquery": kumo.PredictiveQuery,
+    # Note: Table IDs don't have prefixes, so they require explicit
+    # --entity-class
+}
+
+
+def _load_with_class(entity_id: str, entity_class: type) -> object:
+    """Helper to load an entity when the class is known."""
+    # Order of attempts: get_by_name, load, constructor
+    if hasattr(entity_class, "get_by_name"):
+        return entity_class.get_by_name(entity_id)
+    elif hasattr(entity_class, "load"):
+        return entity_class.load(entity_id)
+
+    try:
+        # For jobs like TrainingJob, BatchPredictionJob,
+        # FileUploadConnector
+        return entity_class(entity_id)
+    except (TypeError, AttributeError):
+        pass  # Fall through to the error
+
+    raise NotImplementedError(
+        f"Don't know how to load object of type {entity_class.__name__}")
+
+
+def load_from_id(
+    entity_id: str,
+    entity_class_str: str | None = None,
+) -> object:
+    """Load a Kumo object as an SDK object.
+    - If entity_class_str is provided, it's used to find the type.
+    - If not, the type is inferred from the ID prefix.
+    """
+    try:
+        # Scenario A: Explicit class string provided
+        if entity_class_str:
+            if entity_class_str not in ENTITY_TYPE_MAP:
+                raise ValueError(f"Unknown entity_class '{entity_class_str}'. "
+                                 f"Supported types are: "
+                                 f"{', '.join(ENTITY_TYPE_MAP.keys())}")
+            entity_class = ENTITY_TYPE_MAP[entity_class_str]
+            return _load_with_class(entity_id, entity_class)
+
+        # Scenario B: No class string provided, so infer from ID
+        # prefix
+        prefix = entity_id.split("-", 1)[0].lower()
+        if prefix in ENTITY_PREFIX_MAP:
+            inferred_class = ENTITY_PREFIX_MAP[prefix]
+            return _load_with_class(entity_id, inferred_class)
+        else:
+            raise ValueError(
+                f"Could not infer entity type from ID '{entity_id}'. "
+                "For an entity with a non-prefixed ID"
+                "(like a Connector, Graph, or Table), "
+                "please provide the 'entity_class' parameter. "
+                "Supported prefixes are: " +
+                ", ".join(ENTITY_PREFIX_MAP.keys()) +
+                "\n and supported classes are: " +
+                ", ".join(ENTITY_TYPE_MAP.keys()))
+    except (HTTPException, ValueError) as e:
+        class_name = (entity_class_str
+                      if entity_class_str else "inferred type")
+        raise ValueError(
+            f"Failed to load entity '{entity_id}' of type {class_name}") from e
+    except Exception as e:
+        raise ValueError(f"An unexpected error occurred while"
+                         f"loading entity '{entity_id}'") from e

kumoai/codegen/naming.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import keyword
+from collections import defaultdict
+from typing import Any
+
+from kumoai.codegen.identity import get_config_id
+
+
+def _sanitize_identifier(name: str) -> str:
+    """Sanitize a name to be a valid Python identifier."""
+    if not name:
+        return "obj"
+
+    sanitized = "".join(char if char.isalnum() else "_"
+                        for char in name.lower())
+    sanitized = "_".join(filter(None, sanitized.split("_")))
+
+    if not sanitized:
+        return "obj"
+
+    if sanitized[0].isdigit():
+        sanitized = f"_{sanitized}"
+
+    if keyword.iskeyword(sanitized) or sanitized in dir(__builtins__):
+        sanitized = f"{sanitized}_"
+
+    return sanitized
+
+
+class NameManager:
+    def __init__(self) -> None:
+        self._counts: defaultdict[str, int] = defaultdict(int)
+        self._names: dict[str, str] = {}  # config_id -> variable_name
+
+    def assign_entity_variable(self, obj: Any) -> str:
+        config_id = get_config_id(obj)
+        if config_id in self._names:
+            return self._names[config_id]
+
+        entity_name = ""
+        if hasattr(obj, "name") and obj.name:
+            entity_name = str(obj.name)
+        elif hasattr(obj, "source_name") and obj.source_name:
+            entity_name = str(obj.source_name)
+
+        type_name = obj.__class__.__name__.lower()
+
+        if entity_name:
+            sanitized_name = _sanitize_identifier(entity_name)
+            base_name = (sanitized_name if sanitized_name.replace("_", "")
+                         == type_name else f"{sanitized_name}_{type_name}")
+        else:
+            base_name = type_name
+
+        self._counts[base_name] += 1
+        name = f"{base_name}_{self._counts[base_name]}"
+        self._names[config_id] = name
+        return name
+
+    def assign_temp_variable(self, path: str, value: Any) -> str:
+        base_name = self._get_base_name_for_temp(path, value)
+        self._counts[base_name] += 1
+        return f"{base_name}_{self._counts[base_name]}"
+
+    def _get_base_name_for_temp(self, path: str, value: Any) -> str:
+        if path:
+            if "." in path:
+                parts = path.split(".")
+                for part in reversed(parts):
+                    if part and not part.startswith("["):
+                        return part.split("[")[0]
+            if "[" in path:
+                return path.split("[")[0]
+            if not path.startswith("["):
+                return path
+
+        primitives = (type(None), str, int, float, bool, list, dict, set,
+                      tuple)
+        if not isinstance(value, primitives):
+            import re
+            class_name = type(value).__name__
+            return re.sub(r"(?<!^)(?=[A-Z])", "_", class_name).lower()
+
+        if isinstance(value, list):
+            return "list"
+        if isinstance(value, dict):
+            return "dict"
+        if isinstance(value, set):
+            return "set"
+        if isinstance(value, tuple):
+            return "tuple"
+
+        return "temp_obj"

kumoai/codegen/registry.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import importlib
+import pkgutil
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Type,
+)
+
+from kumoai.codegen.context import CodegenContext
+from kumoai.codegen.edits import UniversalReplacementEdit
+from kumoai.codegen.identity import get_config_id
+from kumoai.codegen.naming import NameManager
+
+
+class Handler(NamedTuple):
+    parents: Callable[[object, CodegenContext],
+                      List[object]]  # Added codegen_ctx parameter
+    required_imports: Callable[[object], List[str]]
+    emit_lines: Callable[[object, str, dict, CodegenContext],
+                         List[str]]  # Added codegen_ctx parameter
+    detect_edits: Optional[Callable[[object, object, NameManager],
+                                    Sequence[UniversalReplacementEdit]]]
+    get_parent_map: Optional[Callable[[object], dict[str, dict[str,
+                                                               Any]]]] = None
+
+
+REG: dict[Type, Handler] = {}
+
+
+def register_shared_parents(ctx: CodegenContext, obj: object,
+                            handler: Handler) -> None:
+    """Register parents that this handler wants
+    to share with other handlers.
+    """
+    if handler.get_parent_map:
+        parent_map = handler.get_parent_map(obj)
+        # parent_map format: {object_id: {key: parent_obj}}
+        for obj_id, shared_data in parent_map.items():
+            ctx.shared_parents[obj_id] = shared_data
+
+
+def lookup_shared_parent(ctx: CodegenContext, obj: object, key: str) -> Any:
+    """Look up a shared parent by key from another handler using config ID."""
+    config_id = get_config_id(obj)
+    return ctx.shared_parents.get(config_id, {}).get(key)
+
+
+def store_shared_parent(ctx: CodegenContext, obj: object, key: str,
+                        parent_obj: object) -> None:
+    """Store a shared parent for an object."""
+    config_id = get_config_id(obj)
+    if config_id not in ctx.shared_parents:
+        ctx.shared_parents[config_id] = {}
+    ctx.shared_parents[config_id][key] = parent_obj
+
+
+def store_object_var(ctx: CodegenContext, obj: object, var_name: str) -> None:
+    """Store the variable name for an object using config_id."""
+    config_id = get_config_id(obj)
+    ctx.object_to_var[config_id] = var_name
+
+
+def get_object_var(ctx: CodegenContext, obj: object) -> str:
+    """Get the variable name for an object using config_id."""
+    config_id = get_config_id(obj)
+    var_name = ctx.object_to_var.get(config_id)
+    if not var_name:
+        raise ValueError(
+            f"No variable name found for object {type(obj).__name__} "
+            f"with config_id {config_id}")
+    return var_name
+
+
+def init_execution_env(ctx: CodegenContext) -> None:
+    """Initialize the execution environment in context."""
+    import kumoai as kumo
+    ctx.execution_env = {"kumo": kumo}
+
+
+def execute_in_env(ctx: CodegenContext, lines: list[str],
+                   imports: Optional[list[str]] = None) -> None:
+    """Execute lines in the context's execution environment."""
+    if imports:
+        for import_line in imports:
+            exec(import_line, ctx.execution_env)
+
+    for line in lines:
+        if line.strip() and not line.strip().startswith("#"):
+            exec(line, ctx.execution_env)
+
+
+def get_from_env(ctx: CodegenContext, var_name: str) -> Any:
+    """Get an object from the context's execution environment."""
+    return ctx.execution_env.get(var_name)
+
+
+def _discover_and_register_handlers() -> None:
+    """Dynamically discover and import all modules in the 'handlers' folders,
+    call their `get_handlers` function, and register the returned handlers.
+    """
+    from . import handlers
+
+    handlers_dir = handlers.__path__
+    prefix = f"{handlers.__name__}."
+
+    for _, module_name, _ in pkgutil.iter_modules(handlers_dir, prefix):
+        module = importlib.import_module(module_name)
+        if hasattr(module, "get_handlers"):
+            handlers_to_register: Dict[Type, Handler] = (module.get_handlers())
+            for cls, handler in handlers_to_register.items():
+                REG[cls] = handler
+
+
+_discover_and_register_handlers()

kumoai/connector/__init__.py

@@ -0,0 +1,31 @@
+from .base import Connector
+from .s3_connector import S3Connector
+from .snowflake_connector import SnowflakeConnector
+from .databricks_connector import DatabricksConnector
+from .bigquery_connector import BigQueryConnector
+from .file_upload_connector import FileUploadConnector
+from .glue_connector import GlueConnector
+from .source_table import (
+    SourceTable,
+    SourceTableFuture,
+    LLMSourceTableFuture,
+    SourceColumn,
+)
+from .utils import upload_table, delete_uploaded_table, replace_table
+
+__all__ = [
+    'Connector',
+    'S3Connector',
+    'SnowflakeConnector',
+    'DatabricksConnector',
+    'BigQueryConnector',
+    'FileUploadConnector',
+    'GlueConnector',
+    'SourceTable',
+    'SourceTableFuture',
+    'LLMSourceTableFuture',
+    'SourceColumn',
+    'upload_table',
+    'delete_uploaded_table',
+    'replace_table',
+]

kumoai/connector/base.py

@@ -0,0 +1,153 @@
+from abc import ABC, abstractmethod
+from typing import List, Set, Union
+
+from kumoapi.data_source import DataSourceType
+from kumoapi.source_table import (
+    BigQuerySourceTableRequest,
+    DatabricksSourceTableRequest,
+    S3SourceTableRequest,
+    SnowflakeSourceTableRequest,
+    SourceTableConfigRequest,
+    SourceTableConfigResponse,
+    SourceTableDataRequest,
+    SourceTableDataResponse,
+    SourceTableListRequest,
+    SourceTableListResponse,
+    SourceTableValidateRequest,
+    SourceTableValidateResponse,
+)
+
+from kumoai import global_state
+from kumoai.connector.source_table import SourceTable
+from kumoai.exceptions import HTTPException
+
+
+class Connector(ABC):
+    r"""A connector to a backing data source, that can be used to create Kumo
+    tables.
+    """
+
+    _validated_tables: Set[str] = set()
+
+    # Metadata ################################################################
+
+    @property
+    @abstractmethod
+    def source_type(self) -> DataSourceType:
+        r"""Returns the data source type accessible by this connector."""
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        r"""Returns the name of the connector.
+
+        .. note::
+            If the connector does not support naming, the name refers to an
+            internal specifier.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _source_table_request(
+        self,
+        table_names: List[str],
+    ) -> Union[S3SourceTableRequest, BigQuerySourceTableRequest,
+               DatabricksSourceTableRequest, SnowflakeSourceTableRequest]:
+        raise NotImplementedError()
+
+    # Tables ##################################################################
+
+    def table_names(self) -> List[str]:
+        r"""Returns a list of table names accessible through this connector."""
+        return self._list_tables().table_names
+
+    def has_table(self, name: str) -> bool:
+        r"""Returns :obj:`True` if the table exists in this connector,
+        :obj:`False` otherwise.
+
+        Args:
+            name: The table name.
+        """
+        try:
+            resp = self._validate_table(name)
+            return resp.is_valid
+        except HTTPException:
+            # In case of HTTPException, Kumo backend doesn't have api
+            # implemented and we skip check by returns True.
+            return True
+
+    def table(self, name: str) -> SourceTable:
+        r"""Returns a :class:`~kumoai.connector.SourceTable` object
+        corresponding to a source table behind this connector. A source table
+        is a view into the raw data of table :obj:`name`. To use a source
+        table in Kumo, you will need to construct a
+        :class:`~kumoai.graph.Table` from the source table.
+
+        Args:
+            name: The table name.
+
+        Raises:
+            :class:`ValueError`: if :obj:`name` does not exist in the backing
+                connector.
+        """
+        if not self.has_table(name):
+            raise ValueError(f"The table '{name}' does not exist in {self}. "
+                             f"Please check the existence of the source data.")
+
+        return SourceTable(name=name, connector=self)
+
+    def _validate_table(self, table_name: str) -> SourceTableValidateResponse:
+        if table_name in self._validated_tables:
+            return SourceTableValidateResponse(is_valid=True, msg='')
+
+        req = SourceTableValidateRequest(
+            connector_id=self.name,
+            table_name=table_name,
+            source_type=self.source_type,
+        )
+        ret = global_state.client.source_table_api.validate_table(req)
+
+        # Cache the result for the whole session.
+        if ret.is_valid:
+            self._validated_tables.add(table_name)
+        return ret
+
+    def _list_tables(self) -> SourceTableListResponse:
+        req = SourceTableListRequest(connector_id=self.name,
+                                     source_type=self.source_type)
+        return global_state.client.source_table_api.list_tables(req)
+
+    def _get_table_data(
+        self,
+        table_names: List[str],
+        sample_rows: int,
+    ) -> SourceTableDataResponse:
+        req = SourceTableDataRequest(
+            source_table_request=self._source_table_request(table_names),
+            sample_rows=sample_rows,
+        )
+        return global_state.client.source_table_api.get_table_data(req)
+
+    def _get_table_config(self, table_name: str) -> SourceTableConfigResponse:
+        # TODO(manan): rest backend for this is a bit broken, it never returns
+        # directories...
+        req = SourceTableConfigRequest(connector_id=self.name,
+                                       table_name=table_name,
+                                       source_type=self.source_type)
+        return global_state.client.source_table_api.get_table_config(req)
+
+    # Class properties ########################################################
+
+    def __hash__(self) -> int:
+        return hash(self.__dict__)
+
+    def __contains__(self, name: str) -> bool:
+        return self.has_table(name)
+
+    def __getitem__(self, name: str) -> SourceTable:
+        return self.table(name)
+
+    def __repr__(self) -> str:
+        # TODO(manan): class-overrideable metadata?
+        return f'{self.__class__.__name__}()'