Used TypeDict for CachingVisitor.state

[TomAugspurger]

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

1. Typos in the key names can cause unexpected KeyErrors
2. The Any type for the values of state means 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.

Now, we'll declare a State TypedDict with some specific fields for each transformer that uses CachingVisitor.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[TomAugspurger]

FYI @rjzamora this will have some overlap with #19130. Any new keys we add to state will need a field in CachingVisitorState.

[TomAugspurger]

Hmm, failing on 3.10: https://github.com/rapidsai/cudf/actions/runs/15594634303/job/43923137155?pr=19135#step:10:7551

 │ │ TypeError: cannot inherit from both a TypedDict type and a non-TypedDict base class

This might need to wait until we drop Python 3.10.

[TomAugspurger]

Just seeing if CI passes with TypedDict from typing_extensions. The docs say this should work. If it does, we can discuss adding that as a dependency for python=3.10. Maybe not worth it.

[TomAugspurger]

I'd like to add a dependency on typing-extensions for just python 3.10. For 3.11 and newer, we don't need typing-extensions. Here's what I think is the desired output:

diff --git a/conda/recipes/cudf-polars/recipe.yaml b/conda/recipes/cudf-polars/recipe.yaml
index 85740e5930..8133c7e8a5 100644
--- a/conda/recipes/cudf-polars/recipe.yaml
+++ b/conda/recipes/cudf-polars/recipe.yaml
@@ -50,6 +50,7 @@ requirements:
     - python
     - pylibcudf =${{ version }}
     - polars >=1.24,<1.31
+    - typing-extensions  # py==310
     - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
   ignore_run_exports:
     by_name:
diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml
index b3d28e027d..9589a8a3f9 100644
--- a/python/cudf_polars/pyproject.toml
+++ b/python/cudf_polars/pyproject.toml
@@ -21,6 +21,7 @@ requires-python = ">=3.10"
 dependencies = [
     "polars>=1.25,<1.31",
     "pylibcudf==25.8.*,>=0.0.0a0",
+    "typing-extensions; python_version < '3.11'",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [
     "Intended Audience :: Developers",

I'd hoped that 4ed5168 would do the trick, but that gives the same output as what's on 25.08 (no change).

[gforsyth]

Hey @TomAugspurger -- I'm not sure if the dependency file generator knows how to handle version-specific dependencies in a pyproject.toml.

I think it's reasonable to do something like:

@@ -834,13 +834,11 @@ dependencies:
         packages:
           - polars>=1.25,<1.31
     specific:
-      - output_types: [conda, requirements, pyproject]
+      - output_types: pyproject
         matrices:
-          - matrix: {py: "3.10"}
-            packages:
-              - typing-extensions
           - matrix: null
-            packages: null
+            packages:
+              - "typing-extensions; python_version < '3.11'"
   run_cudf_polars_experimental:
     common:
       - output_types: [conda, requirements, pyproject]

The other spot where typing_extensions would need to show up is in the conda environment files, but it seems to already be there.

As for the rattler-build recipes, those aren't covered (yet) by dependencies.yaml, so you can go ahead and make the necessary changes directly.

[rjzamora]

If folks are OK with this overall I'll clean up a couple things (docs, where we define these Transformer subtypes) and ping for another round of reviews.

I'll defer to you on whether the extra effort is worth it, and I'll be happy to review.

[wence-]

I think this looks close to ready right now (hence approving).

I think we should be able to get away without explicitly constructing the TypedDict instances (see example).

[TomAugspurger]

/merge

[wence-]

Should have been StateT_co I think

[TomAugspurger]

That errored with

python/cudf_polars/cudf_polars/dsl/traversal.py:128: error: Cannot use a covariant type variable as a parameter  [misc]

Some discussion of this at python/mypy#7049, but I wasn't able to figure out a way around that.

[wence-]

This is because mypy doesn't understand how we're using this state object and doesn't realise that this make_recursive is actually a type constructor rather than a regular function.

We could have fixed this by:

type: ignore (describing why)

defining a copy of CachingVisitor called make_recursive that is identically except for the lack of cache.

[TomAugspurger]

Thanks, I'll fix get that fixed.