Figure.shift_origin: Refactor to support shifting origins temporarily

Author: seisman

pygmt/src/shift_origin.py

@@ -1,21 +1,39 @@
 """
-shift_origin - Shift plot origin in x and/or y directions.
+shift_origin - Shift the plot origin in x and/or y directions.
 """
 
 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
-):
+class shift_origin:  # noqa: N801
     r"""
-    Shift plot origin in x and/or y directions.
+    Shift the plot origin in x and/or y directions.
 
-    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 can be temporary or permanent. If used as a context manager, the shifts
+    are temporary and only apply to the block of code within the context manager. If
+    used as a standalone method, the shifts are permanent and apply to all subsequent
+    plots.
+
+    1.  Use as a context manager to shift the plot origin temporarily:
+
+        .. code-block:: python
+
+            with fig.shift_origin(...):
+                ...  # Other plot commands
+                ...
+
+    2.  Use as a standalone method to shift the plot origin permanently:
+
+        .. code-block:: python
+
+            fig.shift_origin(xshift=12)
+            ...  # Other plot commands
+
+    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 +62,71 @@ def shift_origin(
 
     Examples
     --------
+
     >>> import pygmt
+
+    Shifting the plot origin permanently:
+
+    >>> 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
+    >>> fig.shift_origin(xshift=6)
+    <pygmt.src.shift_origin.shift_origin 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")
+    <pygmt.src.shift_origin.shift_origin 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, 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 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}")
-
-    with Session() as lib:
-        lib.call_module(module="plot", args=args)
+
+    def __init__(
+        self, xshift: float | str | None = None, yshift: float | str | None = None
+    ):
+        """
+        Shift the plot origin in x/y directions and store the shift values.
+        """
+        # self._preprocess()  # pylint: disable=protected-access
+        kwdict = {"T": True, "X": xshift, "Y": yshift}
+        with Session() as lib:
+            lib.call_module(module="plot", args=build_arg_list(kwdict))
+            self._xshift = lib.get_common("X")  # False or xshift in inches
+            self._yshift = lib.get_common("Y")  # False or yshift in inches
+
+    def __enter__(self):
+        """
+        Do nothing but return self.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Revert to the original plot origin if the context manager is exited.
+        """
+        # xshift and yshift are in inches, so we need to negate them to revert to the
+        # original plot origin.
+        kwdict = {
+            "T": True,
+            "X": f"{-1.0 * self._xshift}i" if self._xshift else None,
+            "Y": f"{-1.0 * self._yshift}i" if self._yshift else None,
+        }
+        with Session() as lib:
+            lib.call_module(module="plot", args=build_arg_list(kwdict))

pygmt/tests/test_shift_origin.py

@@ -3,8 +3,8 @@
 """
 
 import pytest
+from pygmt import Figure
 from pygmt.exceptions import GMTInvalidInput
-from pygmt.figure import Figure
 
 
 @pytest.mark.mpl_image_compare
@@ -27,6 +27,28 @@ def test_shift_origin():
     return fig
 
 
+def test_shift_origin_context_manager():
+    """
+    Test Figure.shift_origin.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t1"])
+    fig.shift_origin(xshift=3, yshift=3)
+    fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t2"])
+
+    with fig.shift_origin(xshift=3):
+        fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t3"])
+    fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t4"])
+    with fig.shift_origin(xshift=3):
+        fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t5"])
+    with fig.shift_origin(yshift=3):
+        fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t6"])
+    with fig.shift_origin(xshift=3, yshift=3):
+        fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t7"])
+    fig.basemap(region=[0, 1, 0, 1], projection="X3c", frame=["WSen+t8"])
+    return fig
+
+
 def test_shift_origin_unsupported_xshift_yshift():
     """
     Raise an exception if X/Y/xshift/yshift is used.