Wrap GMT_Put_Strings to pass str columns into GMT C API directly

pygmt/clib/session.py

@@ -44,7 +44,7 @@
     "GMT_IS_SURFACE",
 ]
 
-MODES = ["GMT_CONTAINER_ONLY", "GMT_OUTPUT"]
+MODES = ["GMT_CONTAINER_ONLY", "GMT_IS_OUTPUT"]
 
 REGISTRATIONS = ["GMT_GRID_PIXEL_REG", "GMT_GRID_NODE_REG"]
 
@@ -235,7 +235,7 @@ def __getitem__(self, name):
         value = c_get_enum(session, name.encode())
 
         if value is None or value == -99999:
-            raise GMTCLibError("Constant '{}' doesn't exits in libgmt.".format(name))
+            raise GMTCLibError(f"Constant '{name}' doesn't exist in libgmt.")
 
         return value
 
@@ -511,13 +511,13 @@ def create_data(self, family, geometry, mode, **kwargs):
         ----------
         family : str
             A valid GMT data family name (e.g., ``'GMT_IS_DATASET'``). See the
-            ``data_families`` attribute for valid names.
+            ``FAMILIES`` attribute for valid names.
         geometry : str
             A valid GMT data geometry name (e.g., ``'GMT_IS_POINT'``). See the
-            ``data_geometries`` attribute for valid names.
+            ``GEOMETRIES`` attribute for valid names.
         mode : str
-            A valid GMT data mode (e.g., ``'GMT_OUTPUT'``). See the
-            ``data_modes`` attribute for valid names.
+            A valid GMT data mode (e.g., ``'GMT_IS_OUTPUT'``). See the
+            ``MODES`` attribute for valid names.
         dim : list of 4 integers
             The dimensions of the dataset. See the documentation for the GMT C
             API function ``GMT_Create_Data`` (``src/gmt_api.c``) for the full
@@ -731,7 +731,7 @@ def put_vector(self, dataset, column, vector):
         """
         Attach a numpy 1D array as a column on a GMT dataset.
 
-        Use this functions to attach numpy array data to a GMT dataset and pass
+        Use this function to attach numpy array data to a GMT dataset and pass
         it to GMT modules. Wraps ``GMT_Put_Vector``.
 
         The dataset must be created by :meth:`~gmt.clib.Session.create_data`
@@ -791,11 +791,61 @@ def put_vector(self, dataset, column, vector):
                 )
             )
 
+    def put_strings(self, dataset, family, strings):
+        """
+        Attach a numpy 1D array of dtype str as a column on a GMT dataset.
+
+        Use this function to attach string type numpy array data to a GMT
+        dataset and pass it to GMT modules. Wraps ``GMT_Put_Strings``.
+
+        The dataset must be created by :meth:`~gmt.clib.Session.create_data`
+        first.
+
+        .. warning::
+            The numpy array must be C contiguous in memory. If it comes from a
+            column slice of a 2d array, for example, you will have to make a
+            copy. Use :func:`numpy.ascontiguousarray` to make sure your vector
+            is contiguous (it won't copy if it already is).
+
+        Parameters
+        ----------
+        dataset : :class:`ctypes.c_void_p`
+            The ctypes void pointer to a ``GMT_Dataset``. Create it with
+            :meth:`~gmt.clib.Session.create_data`.
+        family : str
+            The family type of the dataset. Can be either ``GMT_IS_VECTOR`` or
+            ``GMT_IS_MATRIX``.
+        strings : numpy 1d-array
+            The array that will be attached to the dataset. Must be a 1d C
+            contiguous array.
+
+        Raises
+        ------
+        GMTCLibError
+            If given invalid input or ``GMT_Put_Strings`` exits with status !=
+            0.
+
+        """
+        c_put_strings = self.get_libgmt_func(
+            "GMT_Put_Strings",
+            argtypes=[ctp.c_void_p, ctp.c_uint, ctp.c_void_p, ctp.c_void_p],
+            restype=ctp.c_int,
+        )
+
+        strings_pointer = strings.ctypes.data_as(ctp.c_char_p)
+        status = c_put_strings(
+            self.session_pointer, self[family], dataset, strings_pointer
+        )
+        if status != 0:
+            raise GMTCLibError(
+                f"Failed to put strings of type {strings.dtype} into dataset"
+            )
+
     def put_matrix(self, dataset, matrix, pad=0):
         """
         Attach a numpy 2D array to a GMT dataset.
 
-        Use this functions to attach numpy array data to a GMT dataset and pass
+        Use this function to attach numpy array data to a GMT dataset and pass
         it to GMT modules. Wraps ``GMT_Put_Matrix``.
 
         The dataset must be created by :meth:`~gmt.clib.Session.create_data`
@@ -856,12 +906,12 @@ def write_data(self, family, geometry, mode, wesn, output, data):
         ----------
         family : str
             A valid GMT data family name (e.g., ``'GMT_IS_DATASET'``). See the
-            ``data_families`` attribute for valid names. Don't use the
+            ``FAMILIES`` attribute for valid names. Don't use the
             ``GMT_VIA_VECTOR`` or ``GMT_VIA_MATRIX`` constructs for this. Use
             ``GMT_IS_VECTOR`` and ``GMT_IS_MATRIX`` instead.
         geometry : str
             A valid GMT data geometry name (e.g., ``'GMT_IS_POINT'``). See the
-            ``data_geometries`` attribute for valid names.
+            ``GEOMETRIES`` attribute for valid names.
         mode : str
             How the data is to be written to the file. This option varies
             depending on the given family. See the GMT API documentation for

pygmt/tests/test_clib.py

@@ -304,159 +304,6 @@ def test_create_data_fails():
                 )
 
 
-def test_put_vector():
-    "Check that assigning a numpy array to a dataset works"
-    dtypes = "float32 float64 int32 int64 uint32 uint64".split()
-    for dtype in dtypes:
-        with clib.Session() as lib:
-            dataset = lib.create_data(
-                family="GMT_IS_DATASET|GMT_VIA_VECTOR",
-                geometry="GMT_IS_POINT",
-                mode="GMT_CONTAINER_ONLY",
-                dim=[3, 5, 1, 0],  # columns, rows, layers, dtype
-            )
-            x = np.array([1, 2, 3, 4, 5], dtype=dtype)
-            y = np.array([6, 7, 8, 9, 10], dtype=dtype)
-            z = np.array([11, 12, 13, 14, 15], dtype=dtype)
-            lib.put_vector(dataset, column=lib["GMT_X"], vector=x)
-            lib.put_vector(dataset, column=lib["GMT_Y"], vector=y)
-            lib.put_vector(dataset, column=lib["GMT_Z"], vector=z)
-            # Turns out wesn doesn't matter for Datasets
-            wesn = [0] * 6
-            # Save the data to a file to see if it's being accessed correctly
-            with GMTTempFile() as tmp_file:
-                lib.write_data(
-                    "GMT_IS_VECTOR",
-                    "GMT_IS_POINT",
-                    "GMT_WRITE_SET",
-                    wesn,
-                    tmp_file.name,
-                    dataset,
-                )
-                # Load the data and check that it's correct
-                newx, newy, newz = tmp_file.loadtxt(unpack=True, dtype=dtype)
-                npt.assert_allclose(newx, x)
-                npt.assert_allclose(newy, y)
-                npt.assert_allclose(newz, z)
-
-
-def test_put_vector_invalid_dtype():
-    "Check that it fails with an exception for invalid data types"
-    with clib.Session() as lib:
-        dataset = lib.create_data(
-            family="GMT_IS_DATASET|GMT_VIA_VECTOR",
-            geometry="GMT_IS_POINT",
-            mode="GMT_CONTAINER_ONLY",
-            dim=[2, 3, 1, 0],  # columns, rows, layers, dtype
-        )
-        data = np.array([37, 12, 556], dtype="object")
-        with pytest.raises(GMTInvalidInput):
-            lib.put_vector(dataset, column=1, vector=data)
-
-
-def test_put_vector_wrong_column():
-    "Check that it fails with an exception when giving an invalid column"
-    with clib.Session() as lib:
-        dataset = lib.create_data(
-            family="GMT_IS_DATASET|GMT_VIA_VECTOR",
-            geometry="GMT_IS_POINT",
-            mode="GMT_CONTAINER_ONLY",
-            dim=[1, 3, 1, 0],  # columns, rows, layers, dtype
-        )
-        data = np.array([37, 12, 556], dtype="float32")
-        with pytest.raises(GMTCLibError):
-            lib.put_vector(dataset, column=1, vector=data)
-
-
-def test_put_vector_2d_fails():
-    "Check that it fails with an exception for multidimensional arrays"
-    with clib.Session() as lib:
-        dataset = lib.create_data(
-            family="GMT_IS_DATASET|GMT_VIA_VECTOR",
-            geometry="GMT_IS_POINT",
-            mode="GMT_CONTAINER_ONLY",
-            dim=[1, 6, 1, 0],  # columns, rows, layers, dtype
-        )
-        data = np.array([[37, 12, 556], [37, 12, 556]], dtype="int32")
-        with pytest.raises(GMTInvalidInput):
-            lib.put_vector(dataset, column=0, vector=data)
-
-
-def test_put_matrix():
-    "Check that assigning a numpy 2d array to a dataset works"
-    dtypes = "float32 float64 int32 int64 uint32 uint64".split()
-    shape = (3, 4)
-    for dtype in dtypes:
-        with clib.Session() as lib:
-            dataset = lib.create_data(
-                family="GMT_IS_DATASET|GMT_VIA_MATRIX",
-                geometry="GMT_IS_POINT",
-                mode="GMT_CONTAINER_ONLY",
-                dim=[shape[1], shape[0], 1, 0],  # columns, rows, layers, dtype
-            )
-            data = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
-            lib.put_matrix(dataset, matrix=data)
-            # wesn doesn't matter for Datasets
-            wesn = [0] * 6
-            # Save the data to a file to see if it's being accessed correctly
-            with GMTTempFile() as tmp_file:
-                lib.write_data(
-                    "GMT_IS_MATRIX",
-                    "GMT_IS_POINT",
-                    "GMT_WRITE_SET",
-                    wesn,
-                    tmp_file.name,
-                    dataset,
-                )
-                # Load the data and check that it's correct
-                newdata = tmp_file.loadtxt(dtype=dtype)
-                npt.assert_allclose(newdata, data)
-
-
-def test_put_matrix_fails():
-    "Check that put_matrix raises an exception if return code is not zero"
-    # It's hard to make put_matrix fail on the C API level because of all the
-    # checks on input arguments. Mock the C API function just to make sure it
-    # works.
-    with clib.Session() as lib:
-        with mock(lib, "GMT_Put_Matrix", returns=1):
-            with pytest.raises(GMTCLibError):
-                lib.put_matrix(dataset=None, matrix=np.empty((10, 2)), pad=0)
-
-
-def test_put_matrix_grid():
-    "Check that assigning a numpy 2d array to a grid works"
-    dtypes = "float32 float64 int32 int64 uint32 uint64".split()
-    wesn = [10, 15, 30, 40, 0, 0]
-    inc = [1, 1]
-    shape = ((wesn[3] - wesn[2]) // inc[1] + 1, (wesn[1] - wesn[0]) // inc[0] + 1)
-    for dtype in dtypes:
-        with clib.Session() as lib:
-            grid = lib.create_data(
-                family="GMT_IS_GRID|GMT_VIA_MATRIX",
-                geometry="GMT_IS_SURFACE",
-                mode="GMT_CONTAINER_ONLY",
-                ranges=wesn[:4],
-                inc=inc,
-                registration="GMT_GRID_NODE_REG",
-            )
-            data = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
-            lib.put_matrix(grid, matrix=data)
-            # Save the data to a file to see if it's being accessed correctly
-            with GMTTempFile() as tmp_file:
-                lib.write_data(
-                    "GMT_IS_MATRIX",
-                    "GMT_IS_SURFACE",
-                    "GMT_CONTAINER_AND_DATA",
-                    wesn,
-                    tmp_file.name,
-                    grid,
-                )
-                # Load the data and check that it's correct
-                newdata = tmp_file.loadtxt(dtype=dtype)
-                npt.assert_allclose(newdata, data)
-
-
 def test_virtual_file():
     "Test passing in data via a virtual file with a Dataset"
     dtypes = "float32 float64 int32 int64 uint32 uint64".split()