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(