Add the Frame/Axis class for setting frame and axes

Author: seisman

pygmt/params/__init__.py

@@ -3,5 +3,6 @@
 """
 
 from pygmt.params.box import Box
+from pygmt.params.frame import Axis, Frame
 from pygmt.params.pattern import Pattern
 from pygmt.params.position import Position

pygmt/params/frame.py

@@ -0,0 +1,120 @@
+"""
+The Axes, Axis, and Frame classes for specifying the frame.
+"""
+
+import dataclasses
+
+from pygmt.alias import Alias
+from pygmt.exceptions import GMTParameterError
+from pygmt.params.base import BaseParam
+
+
+@dataclasses.dataclass(repr=False)
+class Axis(BaseParam):
+    """
+    Class for setting up one axis of a plot.
+    """
+
+    # Specify the intervals for annotations, ticks, and gridlines.
+    annot: float | None = None
+    tick: float | None = None
+    grid: float | None = None
+    # The label for the axis [Default is no label].
+    label: str | None = None
+
+    @property
+    def _aliases(self):
+        return [
+            Alias(self.annot, name="annot", prefix="a"),
+            Alias(self.tick, name="tick", prefix="f"),
+            Alias(self.grid, name="grid", prefix="g"),
+            Alias(self.label, name="label", prefix="+l"),
+        ]
+
+
+@dataclasses.dataclass(repr=False)
+class _Axes(BaseParam):
+    """
+    A private class to build the Axes part of the Frame class.
+    """
+
+    # Refer to the Frame class for full documentation.
+    axes: str | None = None
+    title: str | None = None
+
+    @property
+    def _aliases(self):
+        return [
+            Alias(self.axes, name="axes"),
+            Alias(self.title, name="title", prefix="+t"),
+        ]
+
+
+@dataclasses.dataclass(repr=False)
+class Frame(BaseParam):
+    """
+    Class for setting up the frame of a plot.
+    """
+
+    #: Specify which axes to draw and their attributes.
+    axes: str | None = None
+
+    #: The title string centered above the plot frame [Default is no title].
+    title: str | None = None
+
+    #: Specify the attributes for axes.
+    #:
+    #: The attributes for each axis can be specified in two ways:
+    #: #. specifying the same attributes for all axes using the ``axis`` parameter
+    #: #. specifying the attributes for each axis separately using the ``xaxis``,
+    #: ``yaxis``, ``zaxis`` parameter for the x-, y, and z-axes, respectively.
+    #:
+    #: GMT uses the notion of primary (the default) and secondary axes, while secondary
+    #: axes are optional and mostly used for time axes annotations. To specify the
+    #: attributes for the secondary axes, use the ``xaxis2``, ``yaxis2``, and ``zaxis2``
+    #: parameters.
+    axis: Axis | None = None
+    xaxis: Axis | None = None
+    yaxis: Axis | None = None
+    zaxis: Axis | None = None
+    xaxis2: Axis | None = None
+    yaxis2: Axis | None = None
+    zaxis2: Axis | None = None
+
+    def _validate(self):
+        """
+        Validate the parameters of the Frame class.
+        """
+        if self.axis is not None and any(
+            [self.xaxis, self.yaxis, self.zaxis, self.xaxis2, self.yaxis2, self.zaxis2]
+        ):
+            raise GMTParameterError(
+                conflicts_with=(
+                    "axis",
+                    ["xaxis", "yaxis", "zaxis", "xaxis2", "yaxis2", "zaxis2"],
+                ),
+                reason="Either 'axis' or the individual axis parameters can be set, but not both.",
+            )
+
+    @property
+    def _aliases(self):
+        return [
+            Alias(_Axes(axes=self.axes, title=self.title)),
+            Alias(self.axis, name="axis"),
+            Alias(self.xaxis, name="xaxis", prefix="px" if self.xaxis2 else "x"),
+            Alias(self.yaxis, name="yaxis", prefix="py" if self.yaxis2 else "y"),
+            Alias(self.zaxis, name="zaxis", prefix="pz" if self.zaxis2 else "z"),
+            Alias(self.xaxis2, name="xaxis2", prefix="sx"),
+            Alias(self.yaxis2, name="yaxis2", prefix="sy"),
+            Alias(self.zaxis2, name="zaxis2", prefix="sz"),
+        ]
+
+    def __iter__(self):
+        """
+        Iterate over the aliases of the class.
+
+        Yields
+        ------
+        The value of each alias in the class. None are excluded.
+        """
+        yield from (alias._value for alias in self._aliases if alias._value is not None)

pygmt/tests/test_params_frame.py

@@ -0,0 +1,88 @@
+"""
+Test the Frame and Axis classes.
+"""
+
+from pygmt.params import Axis, Frame
+
+
+def test_params_axis():
+    """
+    Test the Axis class.
+    """
+    assert str(Axis(annot=True)) == "a"
+    assert str(Axis(annot=True, tick=True, grid=True)) == "afg"
+    assert str(Axis(annot=30, tick=15, grid=5)) == "a30f15g5"
+    assert str(Axis(annot=30, label="LABEL")) == "a30+lLABEL"
+
+
+def test_params_frame_only():
+    """
+    Test the Frame class.
+    """
+    assert str(Frame("WSen")) == "WSen"
+    assert str(Frame(axes="WSEN", title="My Title")) == "WSEN+tMy Title"
+
+
+def test_params_frame_axis():
+    """
+    Test the Frame class with uniform axis setting.
+    """
+    frame = Frame(axes="lrtb", title="My Title", axis=Axis(annot=30, tick=15, grid=10))
+    assert list(frame) == ["lrtb+tMy Title", "a30f15g10"]
+
+    frame = Frame(
+        axes="WSEN",
+        title="My Title",
+        axis=Axis(annot=True, tick=True, grid=True, label="LABEL"),
+    )
+    assert list(frame) == ["WSEN+tMy Title", "afg+lLABEL"]
+
+
+def test_params_frame_separate_axes():
+    """
+    Test the Frame class with separate axis settings.
+    """
+    frame = Frame(
+        axes="lrtb",
+        title="My Title",
+        xaxis=Axis(annot=10, tick=5, grid=2),
+        yaxis=Axis(annot=20, tick=10, grid=4),
+    )
+    assert list(frame) == ["lrtb+tMy Title", "xa10f5g2", "ya20f10g4"]
+
+    frame = Frame(
+        axes="WSEN",
+        title="My Title",
+        xaxis=Axis(annot=True, tick=True, grid=True, label="X-LABEL"),
+        yaxis=Axis(annot=True, tick=True, grid=True, label="Y-LABEL"),
+    )
+    assert list(frame) == ["WSEN+tMy Title", "xafg+lX-LABEL", "yafg+lY-LABEL"]
+
+
+def test_params_frame_separate_axis_secondary():
+    """
+    Test the Frame class with separate axis settings including secondary axes.
+    """
+    frame = Frame(
+        axes="lrtb",
+        title="My Title",
+        xaxis=Axis(annot=10, tick=5, grid=2),
+        xaxis2=Axis(annot=15, tick=7, grid=3),
+        yaxis=Axis(annot=20, tick=10, grid=4),
+        yaxis2=Axis(annot=25, tick=12, grid=5),
+    )
+    assert list(frame) == [
+        "lrtb+tMy Title",
+        "pxa10f5g2",
+        "pya20f10g4",
+        "sxa15f7g3",
+        "sya25f12g5",
+    ]
+
+    frame = Frame(
+        axes="WSEN",
+        title="My Title",
+        xaxis=Axis(annot=True, tick=True, grid=True, label="X-LABEL"),
+        yaxis=Axis(annot=True, tick=True, grid=True, label="Y-LABEL"),
+    )
+    assert list(frame) == ["WSEN+tMy Title", "xafg+lX-LABEL", "yafg+lY-LABEL"]