Figure.clip: Initial implementation

Author: seisman

pygmt/figure.py

@@ -8,6 +8,8 @@
 from tempfile import TemporaryDirectory
 from typing import Literal, overload
 
+from pygmt.src import clip
+
 try:
     import IPython
 
@@ -100,6 +102,8 @@ def __init__(self) -> None:
         self._preview_dir = TemporaryDirectory(prefix=f"{self._name}-preview-")
         self._activate_figure()
 
+        self.clip = clip()
+
     def __del__(self) -> None:
         """
         Clean up the temporary directory that stores the previews.
@@ -411,6 +415,7 @@ def _repr_html_(self) -> str:
 
     from pygmt.src import (  # type: ignore[misc]
         basemap,
+        clip,
         coast,
         colorbar,
         contour,

pygmt/src/__init__.py

@@ -5,6 +5,7 @@
 from pygmt.src.basemap import basemap
 from pygmt.src.binstats import binstats
 from pygmt.src.blockm import blockmean, blockmedian, blockmode
+from pygmt.src.clip import clip
 from pygmt.src.coast import coast
 from pygmt.src.colorbar import colorbar
 from pygmt.src.config import config

pygmt/src/clip.py

@@ -0,0 +1,83 @@
+"""
+Clip.
+"""
+
+from pygmt.clib import Session
+from pygmt.helpers import build_arg_list
+
+
+def _clip_method(clip_type):
+    """
+    Return the clip method for the given clip type.
+    """
+    match clip_type:
+        case "polygon":
+            return "clip"
+        case "land":
+            return "coast"
+        case "water":
+            return "coast"
+
+
+class clip:  # noqa: N801
+    """
+    Clip.
+    """
+
+    def __init__(self):
+        self.type = None
+        self.data = None
+        self.args_enter = ""
+        self.args_exit = ""
+
+    def polygon(self, x, y):
+        """
+        Clip the data to a polygon.
+        """
+        self.type = "polygon"
+        self.data = (x, y)
+        self.args_enter = []
+        self.args_exit = ["-C"]
+        return self
+
+    def land(self):
+        """
+        Clip the data to the land.
+        """
+        self.type = "land"
+        self.data = None
+        self.args_enter = build_arg_list({"G": True})
+        self.args_exit = build_arg_list({"Q": True})
+        return self
+
+    def water(self):
+        """
+        Clip the data to the water.
+        """
+        self.type = "water"
+        self.data = None
+        self.args_enter = build_arg_list({"S": True})
+        self.args_exit = build_arg_list({"Q": True})
+        return self
+
+    def __enter__(self):
+        """
+        Enter the context manager.
+        """
+        module = _clip_method(self.type)
+
+        with Session() as lib:
+            if module == "clip":
+                with lib.virtualfile_in(x=self.data[0], y=self.data[1]) as vintbl:
+                    lib.call_module(module, args=[vintbl])
+            else:
+                lib.call_module(module, args=self.args_enter)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Exit the context manager.
+        """
+        module = _clip_method(self.type)
+        with Session() as lib:
+            lib.call_module(module, args=self.args_exit)