Use make labs a dataclass

This improves IDE autocomplete.

plotnine/labels.py

@@ -1,35 +1,46 @@
 from __future__ import annotations
 
-import typing
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING
 
-from .exceptions import PlotnineError
 from .iapi import labels_view
-from .mapping.aes import SCALED_AESTHETICS, rename_aesthetics
+from .mapping.aes import rename_aesthetics
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     import plotnine as p9
 
 __all__ = ("xlab", "ylab", "labs", "ggtitle")
-VALID_LABELS = SCALED_AESTHETICS | {"caption", "title", "subtitle"}
 
 
+@dataclass
 class labs:
     """
-    Add labels for aesthetics and/or title
-
-    Parameters
-    ----------
-    kwargs :
-        Aesthetics (with scales) to be renamed. You can also
-        set the `title` and `caption`.
+    Add labels for any aesthetics with a scale or title, subtitle & caption
     """
 
-    labels: labels_view
-
-    def __init__(self, **kwargs: str):
-        unknown = kwargs.keys() - VALID_LABELS
-        if unknown:
-            raise PlotnineError(f"Cannot deal with these labels: {unknown}")
+    # Names of Scaled Aesthetics
+    x: str | None = None
+    y: str | None = None
+    alpha: str | None = None
+    color: str | None = None
+    colour: str | None = None
+    fill: str | None = None
+    linetype: str | None = None
+    shape: str | None = None
+    size: str | None = None
+    stroke: str | None = None
+
+    # Other texts
+    title: str | None = None
+    subtitle: str | None = None
+    caption: str | None = None
+
+    def __post_init__(self):
+        kwargs: dict[str, str] = {
+            f.name: value
+            for f in fields(self)
+            if (value := getattr(self, f.name)) is not None
+        }
         self.labels = labels_view(**rename_aesthetics(kwargs))
 
     def __radd__(self, plot: p9.ggplot) -> p9.ggplot:
@@ -42,34 +53,32 @@ def __radd__(self, plot: p9.ggplot) -> p9.ggplot:
 
 class xlab(labs):
     """
-    Create x-axis label
+    Label/name for the x aesthetic
 
     Parameters
     ----------
-    xlab :
-        x-axis label
+    name :
+        x aesthetic label (x-axis)
     """
 
-    def __init__(self, xlab: str):
-        if xlab is None:
-            raise PlotnineError("Arguments to xlab cannot be None")
-        self.labels = labels_view(x=xlab)
+    def __init__(self, name: str, /):
+        self.x = name
+        self.labels = labels_view(x=name)
 
 
 class ylab(labs):
     """
-    Create y-axis label
+    Label/name for the y aesthetic
 
     Parameters
     ----------
-    ylab :
-        y-axis label
+    name :
+        y aesthetic label i.e. y-axis label
     """
 
-    def __init__(self, ylab: str):
-        if ylab is None:
-            raise PlotnineError("Arguments to ylab cannot be None")
-        self.labels = labels_view(y=ylab)
+    def __init__(self, name: str, /):
+        self.y = name
+        self.labels = labels_view(y=name)
 
 
 class ggtitle(labs):
@@ -82,7 +91,6 @@ class ggtitle(labs):
         Plot title
     """
 
-    def __init__(self, title: str):
-        if title is None:
-            raise PlotnineError("Arguments to ggtitle cannot be None")
+    def __init__(self, title: str, /):
+        self.title = title
         self.labels = labels_view(title=title)

tests/test_ggplot_internals.py

@@ -55,18 +55,6 @@ def test_labels():
     assert gg.labels.title == "title2"
     assert gg.labels.caption == "caption2"
 
-    with pytest.raises(PlotnineError):
-        gg = gg + labs(z="z-axis")
-
-    with pytest.raises(PlotnineError):
-        gg = gg + xlab(None)
-
-    with pytest.raises(PlotnineError):
-        gg = gg + ylab(None)
-
-    with pytest.raises(PlotnineError):
-        gg = gg + ggtitle(None)
-
 
 def test_ggplot_parameters():
     p = ggplot(data, aes("x"))