From 756b37f892362120ebe09a89ee5b69c66c7da8ff Mon Sep 17 00:00:00 2001
From: dangotbanned <125183946+dangotbanned@users.noreply.github.com>
Date: Wed, 7 Aug 2024 16:29:13 +0100
Subject: [PATCH] feat: Adds `@register_theme` decorator

Resolves one item in https://github.com/vega/altair/issues/3519
---
 altair/vegalite/v5/theme.py | 92 +++++++++++++++++++++++++++++++++++--
 1 file changed, 89 insertions(+), 3 deletions(-)

diff --git a/altair/vegalite/v5/theme.py b/altair/vegalite/v5/theme.py
index a3c621658..d65f7ea80 100644
--- a/altair/vegalite/v5/theme.py
+++ b/altair/vegalite/v5/theme.py
@@ -2,13 +2,23 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, Literal
+import sys
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, Dict, Final, Literal, TypeVar
+
+if sys.version_info >= (3, 10):
+    from typing import ParamSpec
+else:
+    from typing_extensions import ParamSpec
+
 
 from altair.utils.theme import ThemeRegistry
 
 if TYPE_CHECKING:
-    import sys
-
+    if sys.version_info >= (3, 11):
+        from typing import LiteralString
+    else:
+        from typing_extensions import LiteralString
     if sys.version_info >= (3, 10):
         from typing import TypeAlias
     else:
@@ -53,6 +63,10 @@
 ]
 
 
+P = ParamSpec("P")
+R = TypeVar("R", bound=Dict[str, Any])
+
+
 class VegaTheme:
     """Implementation of a builtin vega theme."""
 
@@ -94,3 +108,75 @@ def __repr__(self) -> str:
     themes.register(theme, VegaTheme(theme))
 
 themes.enable("default")
+
+
+def register_theme(
+    name: LiteralString, *, enable: bool
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """
+    Decorator for registering a theme function.
+
+    Parameters
+    ----------
+    name
+        Unique name assigned in ``alt.themes``.
+    enable
+        Auto-enable the wrapped theme.
+
+    Examples
+    --------
+    Register and enable a theme::
+
+        from __future__ import annotations
+
+        from typing import Any
+        import altair as alt
+
+
+        @alt.register_theme("param_font_size", enable=True)
+        def custom_theme() -> dict[str, Any]:
+            sizes = 12, 14, 16, 18, 20
+            return {
+                "autosize": {"contains": "content", "resize": True},
+                "background": "#F3F2F1",
+                "config": {
+                    "axisX": {"labelFontSize": sizes[1], "titleFontSize": sizes[1]},
+                    "axisY": {"labelFontSize": sizes[1], "titleFontSize": sizes[1]},
+                    "font": "'Lato', 'Segoe UI', Tahoma, Verdana, sans-serif",
+                    "headerColumn": {"labelFontSize": sizes[1]},
+                    "headerFacet": {"labelFontSize": sizes[1]},
+                    "headerRow": {"labelFontSize": sizes[1]},
+                    "legend": {"labelFontSize": sizes[0], "titleFontSize": sizes[1]},
+                    "text": {"fontSize": sizes[0]},
+                    "title": {"fontSize": sizes[-1]},
+                },
+                "height": {"step": 28},
+                "width": 350,
+            }
+
+    Until another theme has been enabled, all charts will use defaults set in ``custom_theme``::
+
+        from vega_datasets import data
+
+        source = data.stocks()
+        lines = (
+            alt.Chart(source, title=alt.Title("Stocks"))
+            .mark_line()
+            .encode(x="date:T", y="price:Q", color="symbol:N")
+        )
+        lines.interactive(bind_y=False)
+
+    """
+
+    def decorate(func: Callable[P, R], /) -> Callable[P, R]:
+        themes.register(name, func)
+        if enable:
+            themes.enable(name)
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorate