Figure.contour: Adjust processing of arguments passed to the "annotation" and "levels" parameters (#2706)

Author: yvonnefroehlich

pygmt/datasets/earth_relief.py

@@ -140,7 +140,7 @@ def load_earth_relief(
     land_only_srtm_resolutions = ["03s", "01s"]
 
     # Map data source to prefix
-    prefix = { 
+    prefix = {
         "igpp": "earth_relief",
         "gebco": "earth_gebco",
         "gebcosi": "earth_gebcosi",

pygmt/src/contour.py

@@ -3,7 +3,13 @@
 """
 
 from pygmt.clib import Session
-from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
+from pygmt.helpers import (
+    build_arg_list,
+    fmt_docstring,
+    is_nonstr_iter,
+    kwargs_to_strings,
+    use_alias,
+)
 
 
 @fmt_docstring
@@ -54,23 +60,26 @@ def contour(self, data=None, x=None, y=None, z=None, **kwargs):
         Arrays of x and y coordinates and values z of the data points.
     {projection}
     {region}
-    annotation : str or int
+    annotation : float, list, or str
         Specify or disable annotated contour levels, modifies annotated
         contours specified in ``levels``.
 
-        - Specify a fixed annotation interval *annot_int* or a
-          single annotation level +\ *annot_int*.
+        - Specify a fixed annotation interval.
+        - Specify a list of annotation levels.
+        - Disable all annotations by setting ``annotation="n"``.
+        - Adjust the appearance by appending different modifiers, e.g.,
+          ``"annot_int+f10p+gred"`` gives annotations with a font size of 10 points and
+          a red filled box. For all available modifiers see :gmt-docs:`contour.html#a`.
     {frame}
-    levels : str or int
+    levels : float, list, or str
         Specify the contour lines to generate.
 
-        - The file name of a CPT file where the color boundaries will
-          be used as contour levels.
-        - The file name of a 2 (or 3) column file containing the contour
-          levels (col 1), (**C**)ontour or (**A**)nnotate (col 2), and optional
-          angle (col 3).
-        - A fixed contour interval *cont_int* or a single contour with
-          +\ *cont_int*.
+        - The file name of a CPT file where the color boundaries will be used as
+          contour levels.
+        - The file name of a 2 (or 3) column file containing the contour levels (col 0),
+          (**C**)ontour or (**A**)nnotate (col 1), and optional angle (col 2).
+        - A fixed contour interval.
+        - A list of contour levels.
     D : str
         Dump contour coordinates.
     E : str
@@ -114,6 +123,17 @@ def contour(self, data=None, x=None, y=None, z=None, **kwargs):
     """
     kwargs = self._preprocess(**kwargs)
 
+    # Specify levels for contours or annotations.
+    # One level is converted to a string with a trailing comma to separate it from
+    # specifying an interval.
+    # Multiple levels are concatenated to a comma-separated string.
+    for arg in ["A", "C"]:
+        if is_nonstr_iter(kwargs.get(arg)):
+            if len(kwargs[arg]) == 1:  # One level
+                kwargs[arg] = str(kwargs[arg][0]) + ","
+            else:  # Multiple levels
+                kwargs[arg] = ",".join(f"{item}" for item in kwargs[arg])
+
     with Session() as lib:
         with lib.virtualfile_in(
             check_kind="vector", data=data, x=x, y=y, z=z, required_z=True

pygmt/tests/baseline/test_contour_interval.png.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: 44d70a0b17bc7c7939462184bf06e4da
+  size: 50998
+  isexec: true
+  hash: md5
+  path: test_contour_interval.png

pygmt/tests/baseline/test_contour_multiple_levels.png.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: 7bef85a616c46b9f05f4dbed07bd703d
+  size: 29247
+  isexec: true
+  hash: md5
+  path: test_contour_multiple_levels.png

pygmt/tests/baseline/test_contour_one_level.png.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: 8c1ed221788e3af76279a7765640ea43
+  size: 28882
+  isexec: true
+  hash: md5
+  path: test_contour_one_level.png

pygmt/tests/test_contour.py

@@ -75,6 +75,60 @@ def test_contour_from_file(region):
     return fig
 
 
+@pytest.mark.mpl_image_compare
+def test_contour_interval(region):
+    """
+    Plot data with fixed (different) contour and annotation intervals.
+    """
+    fig = Figure()
+    fig.contour(
+        data=POINTS_DATA,
+        projection="X10c",
+        region=region,
+        frame="af",
+        levels=0.1,
+        annotation=0.2,
+        pen=True,
+    )
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_contour_one_level(region):
+    """
+    Plot data with one contour level and one (different) annotation level.
+    """
+    fig = Figure()
+    fig.contour(
+        data=POINTS_DATA,
+        projection="X10c",
+        region=region,
+        frame="af",
+        levels=[0.4],
+        annotation=[0.5],
+        pen=True,
+    )
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_contour_multiple_levels(region):
+    """
+    Plot data with multiple (different) contour and annotation levels.
+    """
+    fig = Figure()
+    fig.contour(
+        data=POINTS_DATA,
+        projection="X10c",
+        region=region,
+        frame="af",
+        levels=[0.2, 0.3],
+        annotation=[0.4, 0.45],
+        pen=True,
+    )
+    return fig
+
+
 @pytest.mark.mpl_image_compare(filename="test_contour_vec.png")
 def test_contour_incols_transposed_data(region):
     """