Figure.colorbar: Pythonic way to set labels and annotations

Author: seisman

pygmt/src/colorbar.py

@@ -11,12 +11,58 @@
 from pygmt.exceptions import GMTValueError
 from pygmt.helpers import build_arg_list, fmt_docstring, use_alias
 from pygmt.helpers.utils import is_nonstr_iter
-from pygmt.params import Box, Position
+from pygmt.params import Axis, Box, Frame, Position
 from pygmt.src._common import _parse_position
 
 __doctest_skip__ = ["colorbar"]
 
 
+def _build_frame(
+    annot: float | None = None,
+    tick: float | None = None,
+    grid: float | None = None,
+    annot_angel: float | None = None,
+    annot_prefix: str | None = None,
+    annot_unit: str | None = None,
+    xlabel: str | None = None,
+    ylabel: str | None = None,
+    frame = None,
+):
+    """
+    Create the list of Alias objects for the -B option.
+
+    >>> _build_frame(annot=1, tick=0.5, xlabel="Distance", ylabel="Depth")
+    Frame(xaxis=Axis(annot=1, tick=0.5, label='Distance'), yaxis=Axis(label='Depth'))
+    """
+    if frame is not None and frame is not False:
+        return frame
+
+    _xaxis_is_set = any(
+        v is not None
+        for v in {annot, tick, grid, annot_angel, annot_prefix, annot_unit, xlabel}
+    )
+    _yaxis_is_set = ylabel is not None
+
+    # Need to return None if no parameters are give. Otherwise, it may return "".
+    if not (_xaxis_is_set or _yaxis_is_set):
+        return None
+
+    xaxis, yaxis = None, None
+    if _xaxis_is_set:
+        xaxis = Axis(
+            annot=annot,
+            tick=tick,
+            grid=grid,
+            # angle=annot_angel,
+            # prefix=annot_prefix,
+            # unit=annot_unit,
+            label=xlabel,
+        )
+    if _yaxis_is_set:
+        yaxis = Axis(label=ylabel)
+    return Frame(xaxis=xaxis, yaxis=yaxis)
+
+
 def _alias_option_D(  # noqa: N802, PLR0913
     position=None,
     length=None,
@@ -153,6 +199,14 @@ def colorbar(  # noqa: PLR0913
     length: float | str | None = None,
     width: float | str | None = None,
     orientation: Literal["horizontal", "vertical"] | None = None,
+    annot: float | None = None,
+    tick: float | None = None,
+    grid: float | None = None,
+    annot_angel: float | None = None,
+    annot_prefix: str | None = None,
+    annot_unit: str | None = None,
+    xlabel: str | None = None,
+    ylabel: str | None = None,
     reverse: bool = False,
     nan: bool = False,
     nan_position: Literal["start", "end"] | None = None,
@@ -337,10 +391,9 @@ def colorbar(  # noqa: PLR0913
     >>> fig.basemap(region=[0, 10, 0, 3], projection="X10c/3c", frame=True)
     >>> # Call the colorbar method for the plot
     >>> fig.colorbar(
-    ...     # Set cmap to the "roma" CPT
     ...     cmap="SCM/roma",
-    ...     # Label the x-axis "Velocity" and the y-axis "m/s"
-    ...     frame=["x+lVelocity", "y+lm/s"],
+    ...     xlabel="Velocity",
+    ...     ylabel="m/s",
     ... )
     >>> # Show the plot
     >>> fig.show()
@@ -386,7 +439,17 @@ def colorbar(  # noqa: PLR0913
         Q=Alias(log, name="log"),
         W=Alias(scale, name="scale"),
     ).add_common(
-        B=frame,
+        B=_build_frame(
+            annot=annot,
+            tick=tick,
+            grid=grid,
+            annot_angel=annot_angel,
+            annot_prefix=annot_prefix,
+            annot_unit=annot_unit,
+            xlabel=xlabel,
+            ylabel=ylabel,
+            frame=frame,
+        ),
         J=projection,
         R=region,
         V=verbose,