Figure.shift_origin: Revert back to the original plotting origin using a context manager

pygmt/figure.py

@@ -544,7 +544,6 @@ def _repr_html_(self):
         wiggle,
     )
 
-
 def set_display(method=None):
     """
     Set the display method when calling :meth:`pygmt.Figure.show`.

pygmt/src/shift_origin.py

@@ -2,14 +2,11 @@
 shift_origin - Shift plot origin in x and/or y directions.
 """
 
-from __future__ import annotations
-
 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.
 
@@ -57,12 +54,34 @@
     >>> fig.shift_origin(xshift="w+2c")
     >>> 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):
+        """
+        Exit the context manager.
+        """
+        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/baseline/test_figure_shift_origin.png.dvc → pygmt/tests/baseline/test_shift_origin.png.dvc

@@ -1,4 +1,4 @@
 outs:
 - md5: 85e88836e6d32cded37f792de6439fb9
   size: 9896
-  path: test_figure_shift_origin.png
+  path: test_shift_origin.png

pygmt/tests/test_figure.py

@@ -317,26 +317,6 @@ def test_figure_show():
     fig.show()
 
 
-@pytest.mark.mpl_image_compare
-def test_figure_shift_origin():
-    """
-    Test if fig.shift_origin works.
-    """
-    kwargs = {"region": [0, 3, 0, 5], "projection": "X3c/5c", "frame": 0}
-    fig = Figure()
-    # First call shift_origin without projection and region.
-    # Test issue https://github.com/GenericMappingTools/pygmt/issues/514
-    fig.shift_origin(xshift="2c", yshift="3c")
-    fig.basemap(**kwargs)
-    fig.shift_origin(xshift="4c")
-    fig.basemap(**kwargs)
-    fig.shift_origin(yshift="6c")
-    fig.basemap(**kwargs)
-    fig.shift_origin(xshift="-4c", yshift="6c")
-    fig.basemap(**kwargs)
-    return fig
-
-
 def test_figure_show_invalid_method():
     """
     Test to check if an error is raised when an invalid method is passed to show.

pygmt/tests/test_shift_origin.py

@@ -0,0 +1,26 @@
+"""
+Tests Figure.shift_origin.
+"""
+
+import pytest
+from pygmt import Figure
+
+
+@pytest.mark.mpl_image_compare
+def test_shift_origin():
+    """
+    Test if fig.shift_origin works.
+    """
+    kwargs = {"region": [0, 3, 0, 5], "projection": "X3c/5c", "frame": 0}
+    fig = Figure()
+    # First call shift_origin without projection and region.
+    # Test issue https://github.com/GenericMappingTools/pygmt/issues/514
+    fig.shift_origin(xshift="2c", yshift="3c")
+    fig.basemap(**kwargs)
+    fig.shift_origin(xshift="4c")
+    fig.basemap(**kwargs)
+    fig.shift_origin(yshift="6c")
+    fig.basemap(**kwargs)
+    fig.shift_origin(xshift="-4c", yshift="6c")
+    fig.basemap(**kwargs)
+    return fig