Used TypeDict for CachingVisitor.state

This changes `CachingVisitor.state` to used a `TypedDict`. Previously,
we used a `Mapping[str, Any]`, which had two problems:

1. Risks typos in the key names causing unexpected KeyErrors
2. The `Any` type for the values of `state` mean that, at least by
   default, you won't get any type checking on things using values
   looked up from `state`.

Unit tests should catch all these, but this can provide some earlier
feedback on any issues.

The `total=False` keyword is used in the typeddict to indicate that
these keys are all optional. `mypy` *doesn't* error if you do a lookup
on an optional field of a TypedDict, so we do still need to handle
missing keys like we have been.

Author: TomAugspurger

python/cudf_polars/cudf_polars/dsl/traversal.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """Traversal and visitor utilities for nodes."""
@@ -7,7 +7,7 @@
 
 from typing import TYPE_CHECKING, Any, Generic
 
-from cudf_polars.typing import U_contra, V_co
+from cudf_polars.typing import CachingVisitorState, U_contra, V_co
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator, Mapping, MutableMapping, Sequence
@@ -150,11 +150,11 @@ def __init__(
         self,
         fn: Callable[[U_contra, GenericTransformer[U_contra, V_co]], V_co],
         *,
-        state: Mapping[str, Any] | None = None,
+        state: CachingVisitorState | None = None,
     ) -> None:
         self.fn = fn
         self.cache: MutableMapping[U_contra, V_co] = {}
-        self.state = state if state is not None else {}
+        self.state = state if state is not None else CachingVisitorState()
 
     def __call__(self, value: U_contra) -> V_co:
         """

python/cudf_polars/cudf_polars/experimental/expressions.py

@@ -57,7 +57,7 @@
 
     from cudf_polars.dsl.expressions.base import Expr
     from cudf_polars.dsl.ir import IR
-    from cudf_polars.typing import GenericTransformer, Schema
+    from cudf_polars.typing import CachingVisitorState, GenericTransformer, Schema
     from cudf_polars.utils.config import ConfigOptions
 
 
@@ -510,7 +510,7 @@ def decompose_expr_graph(
     This function recursively decomposes ``named_expr.value`` and
     ``input_ir`` into multiple partition-wise stages.
     """
-    state = {
+    state: CachingVisitorState = {
         "input_ir": input_ir,
         "input_partition_info": partition_info[input_ir],
         "config_options": config_options,

python/cudf_polars/cudf_polars/typing/__init__.py

@@ -21,13 +21,16 @@
 from polars.polars import _expr_nodes as pl_expr, _ir_nodes as pl_ir
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Mapping
+    from collections.abc import Callable, Generator, Mapping
     from typing import TypeAlias
 
     import pylibcudf as plc
 
     from cudf_polars.containers import DataFrame, DataType
     from cudf_polars.dsl import expr, ir, nodebase
+    from cudf_polars.experimental.base import PartitionInfo
+    from cudf_polars.utils.config import ConfigOptions
+
 
 __all__: list[str] = [
     "ClosedInterval",
@@ -158,7 +161,7 @@ def __call__(self, __value: U_contra) -> V_co:
         ...
 
     @property
-    def state(self) -> Mapping[str, Any]:
+    def state(self) -> CachingVisitorState:
         """Arbitrary immutable state."""
         ...
 
@@ -198,3 +201,16 @@ class DataFrameHeader(TypedDict):
 
     columns_kwargs: list[ColumnOptions]
     frame_count: int
+
+
+class CachingVisitorState[NodeT](TypedDict, total=False):
+    """State for CachingVisitor."""
+
+    config_options: ConfigOptions
+    for_parquet: bool
+    input_ir: ir.IR
+    input_partition_info: PartitionInfo
+    name_to_index: Mapping[str, int]
+    replacements: Mapping[NodeT, NodeT]
+    table_ref: plc.expressions.TableReference
+    unique_names: Generator[str, None, None]

python/cudf_polars/docs/overview.md

@@ -371,6 +371,10 @@ def rename(e: Expr, mapping: Mapping[str, str]) -> Expr:
     return mapper(e)
 ```
 
+In practice, `state` is a `TypedDict` defined in `cudf_polars.typing`. To add a
+new field to the state, you'll need to update `CachingVisitorState` to include
+the key and type of that field.
+
 # Containers
 
 Containers should be constructed as relatively lightweight objects

python/cudf_polars/tests/dsl/test_traversal.py

@@ -66,19 +66,23 @@ def test_caching_visitor():
 
     e1 = make_expr(dt, "a", "b")
 
-    mapper = CachingVisitor(rename, state={"mapping": {"b": "c"}})
+    # This test adds an extra key to CachingVisitorState.
+    # That's just a TypedDict, so adding extra keys is fine at runtime.
+    # We don't want to include these test-only fields in the TypeDict,
+    # so we use an ignore here.
+    mapper = CachingVisitor(rename, state={"mapping": {"b": "c"}})  # type: ignore
 
     renamed = mapper(e1)
     assert renamed == make_expr(dt, "a", "c")
     assert len(mapper.cache) == 3
 
     e2 = make_expr(dt, "a", "a")
-    mapper = CachingVisitor(rename, state={"mapping": {"b": "c"}})
+    mapper = CachingVisitor(rename, state={"mapping": {"b": "c"}})  # type: ignore
 
     renamed = mapper(e2)
     assert renamed == make_expr(dt, "a", "a")
     assert len(mapper.cache) == 2
-    mapper = CachingVisitor(rename, state={"mapping": {"a": "c"}})
+    mapper = CachingVisitor(rename, state={"mapping": {"a": "c"}})  # type: ignore
 
     renamed = mapper(e2)
     assert renamed == make_expr(dt, "c", "c")
@@ -189,7 +193,7 @@ def _transform(e: expr.Expr, fn: ExprTransformer) -> expr.Expr:
 
     @_transform.register
     def _(e: expr.Col, fn: ExprTransformer):
-        mapping = fn.state["mapping"]
+        mapping = fn.state["mapping"]  # type: ignore
         if e.name in mapping:
             return type(e)(e.dtype, mapping[e.name])
         return e
@@ -210,7 +214,7 @@ def _rewrite(node: ir.IR, fn: IRTransformer) -> ir.IR:
 
     @_rewrite.register
     def _(node: ir.Select, fn: IRTransformer):
-        expr_mapper = fn.state["expr_mapper"]
+        expr_mapper = fn.state["expr_mapper"]  # type: ignore
         return type(node)(
             node.schema,
             [expr.NamedExpr(e.name, expr_mapper(e.value)) for e in node.exprs],