diff --git a/doc/api/index.rst b/doc/api/index.rst index 75dedbf737d..49743a8de6f 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -216,6 +216,7 @@ Class-style Parameters Axis Box Frame + LegendSpec Pattern Position diff --git a/pygmt/params/__init__.py b/pygmt/params/__init__.py index 43b83ab0131..ebb0e31d368 100644 --- a/pygmt/params/__init__.py +++ b/pygmt/params/__init__.py @@ -4,5 +4,6 @@ from pygmt.params.box import Box from pygmt.params.frame import Axis, Frame +from pygmt.params.legendspec import LegendSpec from pygmt.params.pattern import Pattern from pygmt.params.position import Position diff --git a/pygmt/params/legendspec.py b/pygmt/params/legendspec.py new file mode 100644 index 00000000000..3bf48499bae --- /dev/null +++ b/pygmt/params/legendspec.py @@ -0,0 +1,104 @@ +""" +legendspec - Build a legend specification for Figure.legend. +""" + +import io + + +def _fmt(value: float | str | None) -> str: + """ + Format a value for a legend record. + + >>> _fmt(1.0) + '1.0' + >>> _fmt("1.0c") + '1.0c' + >>> _fmt(None) + '-' + """ + return "-" if value is None else str(value) + + +class LegendSpec: + """ + A legend specification for :meth:`pygmt.Figure.legend`. + """ + + def __init__(self) -> None: + self.records: list[str] = [] + + def _append(self, record: str): + """ + Append one record. + """ + self.records.append(record) + + def add_header(self, text: str, font: str | None = None): + """ + Add a centered header (``H``) [Default font is :gmt-term:`FONT_TITLE`]. + """ + return self._append(f"H {_fmt(font)} {text}") + + def add_symbol( + self, + symbol: str, + size: float | str, + fill: str | None = None, + pen: str | None = None, + label: str | None = None, + dx1: float | str | None = None, + dx2: float | str | None = None, + ): + """ + Add a symbol, with an optional explanatory ``label``. + + ``dx1`` is the offset of the symbol from the left margin of the column and + ``dx2`` the offset of the label; both are computed by GMT if not given. + """ + args = ["S", _fmt(dx1), symbol, f"{size}", _fmt(fill), _fmt(pen)] + if label is not None: + args += [_fmt(dx2), f"{label}"] + self._append(" ".join(args)) + + def add_line( + self, + pen: str | None = None, + length: float | str | None = None, + label: str | None = None, + dx1: float | str | None = None, + dx2: float | str | None = None, + ): + """ + Add a line segment (``S``), with an optional explanatory ``label``. + + A line segment is a symbol record using GMT's horizontal dash symbol, so + ``length`` is the length of the segment and ``pen`` its attributes. + + ``dx1`` is the offset of the segment from the left margin of the column and + ``dx2`` the offset of the label; both are computed by GMT if not given. + """ + self.add_symbol(symbol="-", size=length, pen=pen, label=label, dx1=dx1, dx2=dx2) + + def to_stringio(self) -> io.StringIO: + """ + Return the specification as a :class:`io.StringIO` object. + """ + return io.StringIO(str(self)) + + def __str__(self) -> str: + """ + The legend specification, one record per line. + """ + return "\n".join(self.records) + + def __repr__(self) -> str: + """ + A representation listing the records. + """ + return f"{self.__class__.__name__}({self.records!r})" + + def __len__(self) -> int: + """ + The number of records. + """ + return len(self.records) diff --git a/pygmt/src/legend.py b/pygmt/src/legend.py index 488ca009a17..4bcd0d2bdf1 100644 --- a/pygmt/src/legend.py +++ b/pygmt/src/legend.py @@ -11,14 +11,14 @@ from pygmt.clib import Session from pygmt.exceptions import GMTTypeError from pygmt.helpers import build_arg_list, data_kind, fmt_docstring, is_nonstr_iter -from pygmt.params import Axis, Box, Frame, Position +from pygmt.params import Axis, Box, Frame, LegendSpec, Position from pygmt.src._common import _parse_position @fmt_docstring def legend( self, - spec: PathLike | io.StringIO | None = None, + spec: PathLike | io.StringIO | LegendSpec | None = None, position: Position | Sequence[float | str] | AnchorCode | None = None, width: float | str | None = None, height: float | str | None = None, @@ -71,6 +71,7 @@ def legend( file - Path to the legend specification file - A :class:`io.StringIO` object containing the legend specification + - A :class:`pygmt.params.LegendSpec` object built up record by record See :gmt-docs:`legend.html` for the definition of the legend specification. position @@ -131,6 +132,10 @@ def legend( if height is not None and width is None: width = 0 + # A LegendSpec is passed to GMT as its rendered specification. + if isinstance(spec, LegendSpec): + spec = spec.to_stringio() + kind = data_kind(spec) if kind not in {"empty", "file", "stringio"}: raise GMTTypeError(type(spec))