From c15f549e6aac45cf7c24a7d779897e95f73da805 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Wed, 19 Feb 2025 19:21:42 +0800 Subject: [PATCH 1/5] Figure.shift_origin: Support shifting origins temporarily when used as a context manager --- pygmt/src/shift_origin.py | 98 +++++++++++++++---- .../test_shift_origin_context_manager.png.dvc | 5 + .../test_shift_origin_mixed_modes.png.dvc | 5 + ...hift_origin_nested_context_manager.png.dvc | 5 + pygmt/tests/test_shift_origin.py | 73 +++++++++++++- 5 files changed, 167 insertions(+), 19 deletions(-) create mode 100644 pygmt/tests/baseline/test_shift_origin_context_manager.png.dvc create mode 100644 pygmt/tests/baseline/test_shift_origin_mixed_modes.png.dvc create mode 100644 pygmt/tests/baseline/test_shift_origin_nested_context_manager.png.dvc diff --git a/pygmt/src/shift_origin.py b/pygmt/src/shift_origin.py index f45bdf32559..5b108bc31f3 100644 --- a/pygmt/src/shift_origin.py +++ b/pygmt/src/shift_origin.py @@ -2,20 +2,42 @@ shift_origin - Shift plot origin in x and/or y directions. """ +from contextlib import contextmanager + 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 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 - 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 +66,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) + + >>> 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") + + >>> 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 + + @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() diff --git a/pygmt/tests/baseline/test_shift_origin_context_manager.png.dvc b/pygmt/tests/baseline/test_shift_origin_context_manager.png.dvc new file mode 100644 index 00000000000..201a9a9bc06 --- /dev/null +++ b/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 diff --git a/pygmt/tests/baseline/test_shift_origin_mixed_modes.png.dvc b/pygmt/tests/baseline/test_shift_origin_mixed_modes.png.dvc new file mode 100644 index 00000000000..371a962c77d --- /dev/null +++ b/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 diff --git a/pygmt/tests/baseline/test_shift_origin_nested_context_manager.png.dvc b/pygmt/tests/baseline/test_shift_origin_nested_context_manager.png.dvc new file mode 100644 index 00000000000..e02ffdd7f0b --- /dev/null +++ b/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 diff --git a/pygmt/tests/test_shift_origin.py b/pygmt/tests/test_shift_origin.py index dbbee25ac1e..54ac66a24c4 100644 --- a/pygmt/tests/test_shift_origin.py +++ b/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 shift 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. From 652356c9cae76d8f1078b2f4a2517ca95ef20f48 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Thu, 20 Feb 2025 15:37:28 +0800 Subject: [PATCH 2/5] Update pygmt/tests/test_shift_origin.py Co-authored-by: Michael Grund <23025878+michaelgrund@users.noreply.github.com> --- pygmt/tests/test_shift_origin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pygmt/tests/test_shift_origin.py b/pygmt/tests/test_shift_origin.py index 54ac66a24c4..45b329fcfd2 100644 --- a/pygmt/tests/test_shift_origin.py +++ b/pygmt/tests/test_shift_origin.py @@ -58,7 +58,7 @@ def test_shift_origin_context_manager(): @pytest.mark.mpl_image_compare def test_shift_origin_nested_context_manager(): """ - Test if Figure.shift_origin shift origin correctly when used in a nested context + Test if Figure.shift_origin shifts origin correctly when used in a nested context manager. Expected output is: From 0d489dc55447c2cd95720c466dd6c839209e2078 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Fri, 21 Feb 2025 20:39:45 +0800 Subject: [PATCH 3/5] Update docstrings with permanent shifts first then temporary shifts --- pygmt/src/shift_origin.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/pygmt/src/shift_origin.py b/pygmt/src/shift_origin.py index 5b108bc31f3..738802940c3 100644 --- a/pygmt/src/shift_origin.py +++ b/pygmt/src/shift_origin.py @@ -14,25 +14,24 @@ def shift_origin( r""" Shift the plot origin in x and/or y directions. - 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. + 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 context manager to shift the plot origin temporarily: + 1. Use as a standalone method to shift the plot origin permanently: .. code-block:: python - with fig.shift_origin(...): - ... # Other plot commands - ... + fig.shift_origin(xshift=12) + ... # Other plot commands - 2. Use as a standalone method to shift the plot origin permanently: + 2. Use as a context manager to shift the plot origin temporarily: .. code-block:: python - fig.shift_origin(xshift=12) - ... # Other plot commands + with fig.shift_origin(...): + ... # 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 From 7696cf4f99b846b4d4213264080919fab2f9551e Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Sat, 22 Feb 2025 10:20:54 +0800 Subject: [PATCH 4/5] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Yvonne Fröhlich <94163266+yvonnefroehlich@users.noreply.github.com> --- pygmt/src/shift_origin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pygmt/src/shift_origin.py b/pygmt/src/shift_origin.py index 738802940c3..417e014cb66 100644 --- a/pygmt/src/shift_origin.py +++ b/pygmt/src/shift_origin.py @@ -22,7 +22,7 @@ def shift_origin( .. code-block:: python - fig.shift_origin(xshift=12) + fig.shift_origin(...) ... # Other plot commands 2. Use as a context manager to shift the plot origin temporarily: From a75a440e5106c858f2c6ce60392169ae7b646928 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Sun, 23 Feb 2025 16:43:32 +0800 Subject: [PATCH 5/5] Fix the import style --- pygmt/src/shift_origin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pygmt/src/shift_origin.py b/pygmt/src/shift_origin.py index 417e014cb66..21287847d4d 100644 --- a/pygmt/src/shift_origin.py +++ b/pygmt/src/shift_origin.py @@ -2,7 +2,7 @@ shift_origin - Shift plot origin in x and/or y directions. """ -from contextlib import contextmanager +import contextlib from pygmt.clib import Session from pygmt.helpers import build_arg_list @@ -106,7 +106,7 @@ def shift_origin( _xshift = lib.get_common("X") # False or xshift in inches _yshift = lib.get_common("Y") # False or yshift in inches - @contextmanager + @contextlib.contextmanager def _shift_origin_context(): """ An internal context manager to shift the plot origin temporarily.