chore: move ids to central file (#3839)

Move all IDs to `marimo._types.ids`. will eventually make them more
strongly typed ids too

marimo/_ast/app.py

@@ -26,13 +26,15 @@
 )
 from uuid import uuid4
 
+from marimo._types.ids import CellId_t
+
 if sys.version_info < (3, 10):
     from typing_extensions import ParamSpec, TypeAlias
 else:
     from typing import ParamSpec, TypeAlias
 
 from marimo import _loggers
-from marimo._ast.cell import Cell, CellConfig, CellId_t, CellImpl
+from marimo._ast.cell import Cell, CellConfig, CellImpl
 from marimo._ast.cell_manager import CellManager
 from marimo._ast.errors import (
     CycleError,

marimo/_ast/cell.py

@@ -11,10 +11,9 @@
 from marimo._ast.sql_visitor import SQLVisitor
 from marimo._ast.visitor import ImportData, Language, Name, VariableData
 from marimo._runtime.exceptions import MarimoRuntimeException
+from marimo._types.ids import CellId_t
 from marimo._utils.deep_merge import deep_merge
 
-CellId_t = str
-
 if TYPE_CHECKING:
     from collections.abc import Iterable
     from types import CodeType

marimo/_ast/cell_manager.py

@@ -19,11 +19,12 @@
 else:
     from typing import ParamSpec, TypeAlias
 
-from marimo._ast.cell import Cell, CellConfig, CellId_t
+from marimo._ast.cell import Cell, CellConfig
 from marimo._ast.compiler import cell_factory, toplevel_cell_factory
 from marimo._ast.models import CellData
 from marimo._ast.names import DEFAULT_CELL_NAME
 from marimo._ast.pytest import wrap_fn_for_pytest
+from marimo._types.ids import CellId_t
 
 if TYPE_CHECKING:
     from marimo._ast.app import InternalApp
@@ -75,8 +76,8 @@ def create_cell_id(self) -> CellId_t:
             _id = self.prefix + "".join(
                 self.random_seed.choices(string.ascii_letters, k=4)
             )
-        self.seen_ids.add(_id)
-        return _id
+        self.seen_ids.add(CellId_t(_id))
+        return CellId_t(_id)
 
     # TODO: maybe remove this, it is leaky
     def cell_decorator(

marimo/_ast/compiler.py

@@ -15,13 +15,13 @@
 from marimo import _loggers
 from marimo._ast.cell import (
     Cell,
-    CellId_t,
     CellImpl,
     ImportWorkspace,
     SourcePosition,
 )
 from marimo._ast.variables import is_local
 from marimo._ast.visitor import ImportData, Name, ScopedVisitor
+from marimo._types.ids import CellId_t
 from marimo._utils.tmpdir import get_tmpdir
 
 LOGGER = _loggers.marimo_logger()
@@ -38,7 +38,7 @@ def cell_id_from_filename(filename: str) -> Optional[CellId_t]:
     """Parse cell id from filename."""
     matches = re.findall(r"__marimo__cell_(.*?)_", filename)
     if matches:
-        return str(matches[0])
+        return CellId_t(matches[0])
     return None
 
 

marimo/_ast/models.py

@@ -4,7 +4,8 @@
 from dataclasses import dataclass
 from typing import Optional
 
-from marimo._ast.cell import Cell, CellConfig, CellId_t
+from marimo._ast.cell import Cell, CellConfig
+from marimo._types.ids import CellId_t
 
 
 @dataclass

marimo/_ast/variables.py

@@ -2,16 +2,15 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from marimo._ast.cell import CellId_t
-
-
 from collections import namedtuple
+from typing import Optional
+
+from marimo._types.ids import CellId_t
 
 UnmagledLocal = namedtuple("UnmagledLocal", "name cell")
 
+_EMPTY_CELL_ID = CellId_t("")
+
 
 def if_local_then_mangle(ref: str, cell_id: CellId_t) -> str:
     if is_local(ref):
@@ -21,7 +20,9 @@ def if_local_then_mangle(ref: str, cell_id: CellId_t) -> str:
     return ref
 
 
-def unmangle_local(name: str, cell_id: CellId_t = "") -> UnmagledLocal:
+def unmangle_local(
+    name: str, cell_id: CellId_t = _EMPTY_CELL_ID
+) -> UnmagledLocal:
     if not is_mangled_local(name, cell_id):
         return UnmagledLocal(name, "")
     private_prefix = r"^_cell_\w+?_"
@@ -30,7 +31,7 @@ def unmangle_local(name: str, cell_id: CellId_t = "") -> UnmagledLocal:
     return UnmagledLocal(re.sub(private_prefix, "_", name), name.split("_")[2])
 
 
-def is_mangled_local(name: str, cell_id: CellId_t = "") -> bool:
+def is_mangled_local(name: str, cell_id: CellId_t = _EMPTY_CELL_ID) -> bool:
     return name.startswith(f"_cell_{cell_id}")
 
 
@@ -39,7 +40,7 @@ def is_local(name: str) -> bool:
 
 
 def get_cell_from_local(
-    name: str, cell_id: CellId_t = ""
+    name: str, cell_id: CellId_t = _EMPTY_CELL_ID
 ) -> Optional[CellId_t]:
     local = unmangle_local(if_local_then_mangle(name, cell_id)).cell
     return local if local else None

marimo/_data/models.py

@@ -6,6 +6,8 @@
 from decimal import Decimal
 from typing import Any, List, Literal, Optional, Union
 
+from marimo._types.ids import VariableName
+
 DataType = Literal[
     "string",
     "boolean",
@@ -39,8 +41,6 @@ class DataTableColumn:
 
 DataTableSource = Literal["local", "duckdb", "connection"]
 
-VariableName = str
-
 
 @dataclass
 class DataTable:

marimo/_messaging/console_output_worker.py

@@ -5,9 +5,9 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Literal
 
-from marimo._ast.cell import CellId_t
 from marimo._messaging.cell_output import CellChannel, CellOutput
 from marimo._messaging.mimetypes import KnownMimeType
+from marimo._types.ids import CellId_t
 
 if TYPE_CHECKING:
     from collections import deque

marimo/_messaging/errors.py

@@ -4,8 +4,8 @@
 from dataclasses import dataclass
 from typing import Literal, Optional, Union
 
-from marimo._ast.cell import CellId_t
 from marimo._runtime.dataflow import EdgeWithVar
+from marimo._types.ids import CellId_t
 
 
 @dataclass

marimo/_messaging/ops.py

@@ -27,7 +27,7 @@
 
 from marimo import _loggers as loggers
 from marimo._ast.app import _AppConfig
-from marimo._ast.cell import CellConfig, CellId_t, RuntimeStateType
+from marimo._ast.cell import CellConfig, RuntimeStateType
 from marimo._data.models import (
     ColumnSummary,
     DataSourceConnection,
@@ -54,6 +54,7 @@
 from marimo._runtime.context import get_context
 from marimo._runtime.context.utils import get_mode
 from marimo._runtime.layout.layout import LayoutConfig
+from marimo._types.ids import CellId_t
 from marimo._utils.platform import is_pyodide, is_windows
 
 LOGGER = loggers.marimo_logger()

marimo/_messaging/streams.py

@@ -17,7 +17,6 @@
 )
 
 from marimo import _loggers
-from marimo._ast.cell import CellId_t
 from marimo._messaging.cell_output import CellChannel
 from marimo._messaging.console_output_worker import ConsoleMsg, buffered_writer
 from marimo._messaging.mimetypes import KnownMimeType
@@ -29,6 +28,7 @@
     Stream,
 )
 from marimo._server.types import QueueType
+from marimo._types.ids import CellId_t
 
 if TYPE_CHECKING:
     import queue

marimo/_messaging/types.py

@@ -1,10 +1,12 @@
 # Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
 import abc
 import io
 from typing import Any, Dict, Optional, Tuple
 
-from marimo._ast.cell import CellId_t
 from marimo._messaging.mimetypes import KnownMimeType
+from marimo._types.ids import CellId_t
 
 # The message from the kernel is a tuple of message type
 # and a json representation of the message

marimo/_plugins/ui/_core/registry.py

@@ -13,9 +13,9 @@
     from typing import TypeAlias
 
 from marimo._ast.app import _Namespace
-from marimo._ast.cell import CellId_t
 from marimo._plugins.ui._core.ui_element import UIElement
 from marimo._runtime.context import get_context
+from marimo._types.ids import CellId_t
 
 UIElementId = str
 

marimo/_pyodide/pyodide_session.py

@@ -9,7 +9,7 @@
 from typing import Any, Callable, Optional
 
 from marimo import _loggers
-from marimo._ast.cell import CellConfig, CellId_t
+from marimo._ast.cell import CellConfig
 from marimo._config.config import MarimoConfig
 from marimo._messaging.types import KernelMessage
 from marimo._pyodide.streams import (
@@ -62,6 +62,7 @@
 )
 from marimo._server.session.session_view import SessionView
 from marimo._snippets.snippets import read_snippets
+from marimo._types.ids import CellId_t
 from marimo._utils.case import deep_to_camel_case
 from marimo._utils.formatter import DefaultFormatter
 from marimo._utils.parse_dataclass import parse_raw

marimo/_pyodide/streams.py

@@ -6,7 +6,6 @@
 from typing import Any, Callable, Iterable, Optional
 
 from marimo import _loggers
-from marimo._ast.cell import CellId_t
 from marimo._messaging.cell_output import CellChannel, CellOutput
 from marimo._messaging.mimetypes import KnownMimeType
 from marimo._messaging.ops import CellOp
@@ -18,6 +17,7 @@
     Stdout,
     Stream,
 )
+from marimo._types.ids import CellId_t
 
 LOGGER = _loggers.marimo_logger()
 

marimo/_runtime/app/common.py

@@ -3,7 +3,7 @@
 
 from typing import Any, Dict, Tuple
 
-from marimo._ast.cell import CellId_t
+from marimo._types.ids import CellId_t
 
 OutputsType = Dict[CellId_t, Any]
 DefsType = Dict[str, Any]

marimo/_runtime/app/kernel_runner.py

@@ -4,7 +4,7 @@
 import weakref
 from typing import TYPE_CHECKING, Any
 
-from marimo._ast.cell import CellId_t, CellImpl
+from marimo._ast.cell import CellImpl
 from marimo._config.config import DEFAULT_CONFIG
 from marimo._runtime.app.common import RunOutput
 from marimo._runtime.context.types import get_context
@@ -17,6 +17,7 @@
 )
 from marimo._runtime.runner import cell_runner
 from marimo._server.model import SessionMode
+from marimo._types.ids import CellId_t
 
 if TYPE_CHECKING:
     from marimo._ast.app import InternalApp
@@ -38,7 +39,7 @@ def __init__(self, app: InternalApp) -> None:
         from marimo._runtime.runtime import Kernel
 
         self.app = app
-        self._outputs: dict[str, Any] = {}
+        self._outputs: dict[CellId_t, Any] = {}
 
         ctx = get_context()
         if not isinstance(ctx, KernelRuntimeContext):

marimo/_runtime/app/script_runner.py

@@ -4,7 +4,7 @@
 import asyncio
 from typing import TYPE_CHECKING, Any, Callable, Iterator
 
-from marimo._ast.cell import CellId_t, CellImpl
+from marimo._ast.cell import CellImpl
 from marimo._dependencies.dependencies import DependencyManager
 from marimo._messaging.types import NoopStream
 from marimo._runtime.app.common import RunOutput
@@ -25,6 +25,7 @@
     create_main_module,
     patch_main_module_context,
 )
+from marimo._types.ids import CellId_t
 
 if TYPE_CHECKING:
     from marimo._ast.app import InternalApp

marimo/_runtime/cell_lifecycle_registry.py

@@ -3,8 +3,8 @@
 
 import dataclasses
 
-from marimo._ast.cell import CellId_t
 from marimo._runtime.cell_lifecycle_item import CellLifecycleItem
+from marimo._types.ids import CellId_t
 
 
 @dataclasses.dataclass

marimo/_runtime/context/kernel_context.py

@@ -22,10 +22,10 @@
 
 if TYPE_CHECKING:
     from marimo._ast.app import InternalApp
-    from marimo._ast.cell import CellId_t
     from marimo._messaging.types import Stream
     from marimo._runtime.runtime import Kernel
     from marimo._runtime.state import State
+    from marimo._types.ids import CellId_t
 
 
 @dataclass

marimo/_runtime/context/script_context.py

@@ -29,8 +29,8 @@
 
 if TYPE_CHECKING:
     from marimo._ast.app import InternalApp
-    from marimo._ast.cell import CellId_t
     from marimo._messaging.types import Stream
+    from marimo._types.ids import CellId_t
 
 
 @dataclass

marimo/_runtime/context/types.py

@@ -22,13 +22,13 @@
 
 if TYPE_CHECKING:
     from marimo._ast.app import AppKernelRunnerRegistry, InternalApp
-    from marimo._ast.cell import CellId_t
     from marimo._messaging.types import Stream
     from marimo._output.hypertext import Html
     from marimo._plugins.ui._core.registry import UIElementRegistry
     from marimo._runtime.params import CLIArgs, QueryParams
     from marimo._runtime.state import State, StateRegistry
     from marimo._runtime.virtual_file import VirtualFileRegistry
+    from marimo._types.ids import CellId_t
 
 
 class GlobalContext:

marimo/_runtime/dataflow.py

@@ -7,13 +7,13 @@
 
 from marimo import _loggers
 from marimo._ast.cell import (
-    CellId_t,
     CellImpl,
 )
 from marimo._ast.compiler import code_key
 from marimo._ast.variables import is_mangled_local
 from marimo._ast.visitor import ImportData, Name, VariableData
 from marimo._runtime.executor import execute_cell, execute_cell_async
+from marimo._types.ids import CellId_t
 
 if TYPE_CHECKING:
     from collections.abc import Collection

marimo/_runtime/functions.py

@@ -6,8 +6,8 @@
 import dataclasses
 from typing import Any, Callable, Coroutine, Generic, Type, TypeVar
 
-from marimo._ast.cell import CellId_t
 from marimo._loggers import marimo_logger
+from marimo._types.ids import CellId_t
 from marimo._utils.parse_dataclass import parse_raw
 
 LOGGER = marimo_logger()