From c81d3d8ca621efc96e8c2005c0157938ca099015 Mon Sep 17 00:00:00 2001
From: dangotbanned <125183946+dangotbanned@users.noreply.github.com>
Date: Wed, 7 Aug 2024 19:06:30 +0100
Subject: [PATCH] feat(typing): Generate `VegaThemes` alias

Resolves the *long-term* solution mentioned in https://github.com/vega/altair/issues/3519#issue-2447214476
---
 altair/utils/theme.py                |  4 +--
 altair/vegalite/v5/schema/_typing.py | 17 +++++++++++
 altair/vegalite/v5/theme.py          | 42 +++-------------------------
 tools/generate_schema_wrapper.py     | 17 +++++++----
 4 files changed, 34 insertions(+), 46 deletions(-)

diff --git a/altair/utils/theme.py b/altair/utils/theme.py
index 51929a5fc..02372e690 100644
--- a/altair/utils/theme.py
+++ b/altair/utils/theme.py
@@ -14,14 +14,14 @@
 
 if TYPE_CHECKING:
     from altair.utils.plugin_registry import PluginEnabler
-    from altair.vegalite.v5.theme import _ThemeName
+    from altair.vegalite.v5.theme import AltairThemes, VegaThemes
 
 ThemeType = Callable[..., dict]
 
 
 class ThemeRegistry(PluginRegistry[ThemeType, dict]):
     def enable(
-        self, name: LiteralString | _ThemeName | None = None, **options
+        self, name: LiteralString | AltairThemes | VegaThemes | None = None, **options
     ) -> PluginEnabler:
         """
         Enable a theme by name.
diff --git a/altair/vegalite/v5/schema/_typing.py b/altair/vegalite/v5/schema/_typing.py
index 92ffcb794..bcf6d5090 100644
--- a/altair/vegalite/v5/schema/_typing.py
+++ b/altair/vegalite/v5/schema/_typing.py
@@ -56,6 +56,7 @@
     "TitleOrient_T",
     "TypeForShape_T",
     "Type_T",
+    "VegaThemes",
     "WindowOnlyOp_T",
 ]
 
@@ -76,6 +77,22 @@ def func(
     ): ...
 """
 
+VegaThemes: TypeAlias = Literal[
+    "carbong10",
+    "carbong100",
+    "carbong90",
+    "carbonwhite",
+    "dark",
+    "excel",
+    "fivethirtyeight",
+    "ggplot2",
+    "googlecharts",
+    "latimes",
+    "powerbi",
+    "quartz",
+    "urbaninstitute",
+    "vox",
+]
 Map: TypeAlias = Mapping[str, Any]
 AggregateOp_T: TypeAlias = Literal[
     "argmax",
diff --git a/altair/vegalite/v5/theme.py b/altair/vegalite/v5/theme.py
index a3c621658..c98826826 100644
--- a/altair/vegalite/v5/theme.py
+++ b/altair/vegalite/v5/theme.py
@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, Literal
+from typing import TYPE_CHECKING, Final, Literal, get_args
 
 from altair.utils.theme import ThemeRegistry
+from altair.vegalite.v5.schema._typing import VegaThemes
 
 if TYPE_CHECKING:
     import sys
@@ -14,43 +15,8 @@
     else:
         from typing_extensions import TypeAlias
 
-    # If you add a theme here, also add it in `VEGA_THEMES` below.
-    _ThemeName: TypeAlias = Literal[
-        "default",
-        "carbonwhite",
-        "carbong10",
-        "carbong90",
-        "carbong100",
-        "dark",
-        "excel",
-        "fivethirtyeight",
-        "ggplot2",
-        "googlecharts",
-        "latimes",
-        "opaque",
-        "powerbi",
-        "quartz",
-        "urbaninstitute",
-        "vox",
-    ]
-
-# If you add a theme here, also add it in `_ThemeName` above.
-VEGA_THEMES = [
-    "carbonwhite",
-    "carbong10",
-    "carbong90",
-    "carbong100",
-    "dark",
-    "excel",
-    "fivethirtyeight",
-    "ggplot2",
-    "googlecharts",
-    "latimes",
-    "powerbi",
-    "quartz",
-    "urbaninstitute",
-    "vox",
-]
+AltairThemes: TypeAlias = Literal["default", "opaque"]
+VEGA_THEMES: list[str] = list(get_args(VegaThemes))
 
 
 class VegaTheme:
diff --git a/tools/generate_schema_wrapper.py b/tools/generate_schema_wrapper.py
index 60127ceb6..8ce4b1eb0 100644
--- a/tools/generate_schema_wrapper.py
+++ b/tools/generate_schema_wrapper.py
@@ -19,6 +19,7 @@
 sys.path.insert(0, str(Path.cwd()))
 from tools.schemapi import CodeSnippet, SchemaInfo, codegen
 from tools.schemapi.utils import (
+    TypeAliasTracer,
     get_valid_identifier,
     indent_docstring,
     resolve_references,
@@ -26,6 +27,7 @@
     rst_syntax_for_class,
     ruff_format_py,
     ruff_write_lint_format_str,
+    spell_literal,
 )
 
 SCHEMA_VERSION: Final = "v5.19.0"
@@ -362,9 +364,13 @@ def download_schemafile(
 
 
 def update_vega_themes(fp: Path, /, indent: str | int | None = 2) -> None:
-    data = json.dumps(vlc.get_themes(), indent=indent, sort_keys=True)
+    themes = vlc.get_themes()
+    data = json.dumps(themes, indent=indent, sort_keys=True)
     fp.write_text(data, encoding="utf8")
 
+    theme_names = sorted(iter(themes))
+    TypeAliasTracer.update_aliases(("VegaThemes", spell_literal(theme_names)))
+
 
 def load_schema_with_shorthand_properties(schemapath: Path) -> dict:
     with schemapath.open(encoding="utf8") as f:
@@ -793,6 +799,10 @@ def vegalite_main(skip_download: bool = False) -> None:
         skip_download=skip_download,
     )
 
+    fp_themes = schemapath / "vega-themes.json"
+    print(f"Updating themes\n {schemafile!s}\n  ->{fp_themes!s}")
+    update_vega_themes(fp_themes)
+
     # Generate __init__.py file
     outfile = schemapath / "__init__.py"
     print(f"Writing {outfile!s}")
@@ -850,8 +860,6 @@ def vegalite_main(skip_download: bool = False) -> None:
     files[fp_mixins] = content_mixins
 
     # Write `_typing.py` TypeAlias, for import in generated modules
-    from tools.schemapi.utils import TypeAliasTracer
-
     fp_typing = schemapath / "_typing.py"
     msg = (
         f"Generating\n {schemafile!s}\n  ->{fp_typing!s}\n"
@@ -866,9 +874,6 @@ def vegalite_main(skip_download: bool = False) -> None:
     for fp, contents in files.items():
         print(f"Writing\n {schemafile!s}\n  ->{fp!s}")
         ruff_write_lint_format_str(fp, contents)
-    fp_themes = schemapath / "vega-themes.json"
-    print(f"Updating themes\n {schemafile!s}\n  ->{fp_themes!s}")
-    update_vega_themes(fp_themes)
 
 
 def _create_encode_signature(