Figure.shift_origin: Support shifting origins temporarily when used as a context manager

pygmt/src/shift_origin.py

@@ -2,20 +2,41 @@
 shift_origin - Shift plot origin in x and/or y directions.
 """
 
+import contextlib
+
 from pygmt.clib import Session
+from pygmt.helpers import build_arg_list
 
 
 def shift_origin(
     self, xshift: float | str | None = None, yshift: float | str | None = None
 ):
     r"""
-    Shift plot origin in x and/or y directions.
+    Shift the plot origin in x and/or y directions.
+
+    The shifts can be permanent or temporary. If used as a standalone method, the shifts
+    are permanent and apply to all subsequent plots. If used as a context manager, the
+    shifts are temporary and only apply to the block of code within the context manager.
+
+    1.  Use as a standalone method to shift the plot origin permanently:
+
+        .. code-block:: python
+
+            fig.shift_origin(...)
+            ...  # Other plot commands
+
+    2.  Use as a context manager to shift the plot origin temporarily:
+
+        .. code-block:: python
+
+            with fig.shift_origin(...):
+                ...  # Other plot commands
+                ...
 
-    This method shifts the plot origin relative to the current origin by *xshift* and
-    *yshift* in x and y directions, respectively. Optionally, append the length unit
+    The shifts *xshift* and *yshift* in x and y directions are relative to the current
+    plot origin. The default unit for shifts is centimeters (**c**) but can be changed
+    to other units via :gmt-term:`PROJ_LENGTH_UNIT`. Optionally, append the length unit
     (**c** for centimeters, **i** for inches, or **p** for points) to the shifts.
-    Default unit if not explicitly given is **c**, but can be changed to other units via
-    :gmt-term:`PROJ_LENGTH_UNIT`.
 
     For *xshift*, a special character **w** can also be used, which represents the
     bounding box **width** of the previous plot. The full syntax is
@@ -44,23 +65,63 @@ def shift_origin(
 
     Examples
     --------
+
+    Shifting the plot origin permanently:
+
     >>> import pygmt
     >>> fig = pygmt.Figure()
-    >>> fig.basemap(region=[0, 10, 0, 10], projection="X10c/10c", frame=True)
-    >>> # Shift the plot origin in x direction by 12 cm
-    >>> fig.shift_origin(xshift=12)
-    >>> fig.basemap(region=[0, 10, 0, 10], projection="X14c/10c", frame=True)
-    >>> # Shift the plot origin in x direction based on the previous plot width
-    >>> # Here, the width is 14 cm, and xshift is 16 cm
-    >>> fig.shift_origin(xshift="w+2c")
+    >>> fig.basemap(region=[0, 5, 0, 5], projection="X5c/5c", frame=True)
+    >>> # Shift the plot origin in x direction by 6 cm
+    >>> fig.shift_origin(xshift=6)
+    <contextlib._GeneratorContextManager object at ...>
+    >>> fig.basemap(region=[0, 7, 0, 5], projection="X7c/5c", frame=True)
+    >>> # Shift the plot origin in x direction based on the previous plot width.
+    >>> # Here, the width is 7 cm, and xshift is 8 cm.
+    >>> fig.shift_origin(xshift="w+1c")
+    <contextlib._GeneratorContextManager object at ...>
+    >>> fig.basemap(region=[0, 10, 0, 5], projection="X10c/5c", frame=True)
+    >>> fig.show()
+
+    Shifting the plot origin temporarily:
+
+    >>> fig = pygmt.Figure()
+    >>> fig.basemap(region=[0, 5, 0, 5], projection="X5c/5c", frame=True)
+    >>> # Shift the plot origin in x direction by 6 cm temporarily. The plot origin will
+    >>> # revert back to the original plot origin after the block of code is executed.
+    >>> with fig.shift_origin(xshift=6):
+    ...     fig.basemap(region=[0, 5, 0, 5], projection="X5c/5c", frame=True)
+    >>> # Shift the plot origin in y direction by 6 cm temporarily.
+    >>> with fig.shift_origin(yshift=6):
+    ...     fig.basemap(region=[0, 5, 0, 5], projection="X5c/5c", frame=True)
+    >>> # Shift the plot origin in x and y directions by 6 cm temporarily.
+    >>> with fig.shift_origin(xshift=6, yshift=6):
+    ...     fig.basemap(region=[0, 5, 0, 5], projection="X5c/5c", frame=True)
     >>> fig.show()
     """
     self._preprocess()
-    args = ["-T"]
-    if xshift:
-        args.append(f"-X{xshift}")
-    if yshift:
-        args.append(f"-Y{yshift}")
+    kwdict = {"T": True, "X": xshift, "Y": yshift}
 
     with Session() as lib:
-        lib.call_module(module="plot", args=args)
+        lib.call_module(module="plot", args=build_arg_list(kwdict))
+        _xshift = lib.get_common("X")  # False or xshift in inches
+        _yshift = lib.get_common("Y")  # False or yshift in inches
+
+    @contextlib.contextmanager
+    def _shift_origin_context():
+        """
+        An internal context manager to shift the plot origin temporarily.
+        """
+        try:
+            yield
+        finally:
+            # Revert the plot origin to the original plot origin by shifting it by
+            # -xshift and -yshift in inches.
+            kwdict = {
+                "T": True,
+                "X": f"{-1.0 * _xshift}i" if _xshift else None,
+                "Y": f"{-1.0 * _yshift}i" if _yshift else None,
+            }
+            with Session() as lib:
+                lib.call_module(module="plot", args=build_arg_list(kwdict))
+
+    return _shift_origin_context()

pygmt/tests/baseline/test_shift_origin_context_manager.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 6c63b9935618f607fe22608f7561be22
+  size: 4869
+  hash: md5
+  path: test_shift_origin_context_manager.png

pygmt/tests/baseline/test_shift_origin_mixed_modes.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 3a31df374874ae00920ada0b311b4266
+  size: 5678
+  hash: md5
+  path: test_shift_origin_mixed_modes.png

pygmt/tests/baseline/test_shift_origin_nested_context_manager.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: cd83745f2657ff7daeaba368143db72f
+  size: 4876
+  hash: md5
+  path: test_shift_origin_nested_context_manager.png

pygmt/tests/test_shift_origin.py

@@ -3,8 +3,16 @@
 """
 
 import pytest
+from pygmt import Figure
 from pygmt.exceptions import GMTInvalidInput
-from pygmt.figure import Figure
+
+
+def _numbered_basemap(fig, number, size=3):
+    """
+    A utility function to create a basemap with a number in the center.
+    """
+    fig.basemap(region=[0, 1, 0, 1], projection=f"X{size}c", frame=0)
+    fig.text(position="MC", text=number, font="24p")
 
 
 @pytest.mark.mpl_image_compare
@@ -27,6 +35,69 @@ def test_shift_origin():
     return fig
 
 
+@pytest.mark.mpl_image_compare
+def test_shift_origin_context_manager():
+    """
+    Test if Figure.shift_origin as a context manager shifts origin temporarily.
+
+    Expected output is:
+    | 3 | 4 |
+    | 1 | 2 |
+    """
+    fig = Figure()
+    _numbered_basemap(fig, 1, size=2.5)
+    with fig.shift_origin(xshift=3):
+        _numbered_basemap(fig, 2, size=2.5)
+    with fig.shift_origin(yshift=3):
+        _numbered_basemap(fig, 3, size=2.5)
+    with fig.shift_origin(xshift=3, yshift=3):
+        _numbered_basemap(fig, 4, size=2.5)
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_shift_origin_nested_context_manager():
+    """
+    Test if Figure.shift_origin shifts origin correctly when used in a nested context
+    manager.
+
+    Expected output is:
+    | 4 | 3 |
+    | 1 | 2 |
+    """
+    fig = Figure()
+    _numbered_basemap(fig, 1, size=2.5)
+    with fig.shift_origin(xshift=3):
+        _numbered_basemap(fig, 2, size=2.5)
+        with fig.shift_origin(yshift=3):
+            _numbered_basemap(fig, 3, size=2.5)
+    with fig.shift_origin(yshift=3):
+        _numbered_basemap(fig, 4, size=2.5)
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_shift_origin_mixed_modes():
+    """
+    Test if Figure.shift_origin works when used as a context manager and as a
+    method at the same time.
+
+    Expected output is:
+    |   | 3 | 4 |
+    | 1 | 2 |   |
+    """
+    fig = Figure()
+    _numbered_basemap(fig, 1, size=2.5)
+    with fig.shift_origin(xshift=3):
+        _numbered_basemap(fig, 2, size=2.5)
+    fig.shift_origin(xshift=3)
+    with fig.shift_origin(yshift=3):
+        _numbered_basemap(fig, 3, size=2.5)
+    fig.shift_origin(xshift=3, yshift=3)
+    _numbered_basemap(fig, 4, size=2.5)
+    return fig
+
+
 def test_shift_origin_unsupported_xshift_yshift():
     """
     Raise an exception if X/Y/xshift/yshift is used.