Refactor info and grdinfo to use virtualfile_from_data (#961)

Create a universal `virtualfile_from_data` function
that can handle any raster or vector data input.
This allows us to centralize the data validation
logic in a single place, resulting in a cleaner API
for PyGMT modules to handle different PyData
types (e.g. numpy/pandas/xarray/etc) seamlessly.
As a start, both `info` and `grdinfo` have been
refactored to use this new convenience function.

* Move check_kind to be the first parameter
* Move check_kind out of data_kind function and into virtualfile_from_data

Fixes all the tests which expect first input into data_kind to be the 'data' argument.

* Add doctest example usage of virtualfile_from_data with xarray.Dataset
* Add an entry to doc/api/index.rst for clib.Session.virtualfile_from_data

Author: weiji14

doc/api/index.rst

@@ -186,14 +186,15 @@ the :meth:`~pygmt.clib.Session.call_module` method:
 
     clib.Session.call_module
 
-Passing memory blocks between Python variables (:class:`numpy.ndarray`,
-:class:`pandas.Series`, and :class:`xarray.DataArray`) and GMT happens through *virtual
-files*. These methods are context managers that automate the conversion of Python
-variables to GMT virtual files:
+Passing memory blocks between Python data objects (e.g. :class:`numpy.ndarray`,
+:class:`pandas.Series`, :class:`xarray.DataArray`, etc) and GMT happens through
+*virtual files*. These methods are context managers that automate the
+conversion of Python variables to GMT virtual files:
 
 .. autosummary::
     :toctree: generated
 
+    clib.Session.virtualfile_from_data
     clib.Session.virtualfile_from_matrix
     clib.Session.virtualfile_from_vectors
     clib.Session.virtualfile_from_grid

pygmt/clib/session.py

@@ -25,6 +25,7 @@
     GMTInvalidInput,
     GMTVersionError,
 )
+from pygmt.helpers import data_kind, dummy_context
 
 FAMILIES = [
     "GMT_IS_DATASET",
@@ -1359,6 +1360,90 @@ def virtualfile_from_grid(self, grid):
         with self.open_virtual_file(*args) as vfile:
             yield vfile
 
+    def virtualfile_from_data(self, check_kind=None, data=None, x=None, y=None, z=None):
+        """
+        Store any data inside a virtual file.
+
+        This convenience function automatically detects the kind of data passed
+        into it, and produces a virtualfile that can be passed into GMT later
+        on.
+
+        Parameters
+        ----------
+        check_kind : str
+            Used to validate the type of data that can be passed in. Choose
+            from 'raster', 'vector' or None. Default is None (no validation).
+        data : str, xarray.DataArray, 2d array, or None
+            Any raster or vector data format. This could be a file name, a
+            raster grid, a vector matrix/arrays, or other supported data input.
+        x/y/z : 1d arrays or None
+            x, y and z columns as numpy arrays.
+
+        Returns
+        -------
+        file_context : contextlib._GeneratorContextManager
+            The virtual file stored inside a context manager. Access the file
+            name of this virtualfile using ``with file_context as fname: ...``.
+
+        Examples
+        --------
+        >>> from pygmt.helpers import GMTTempFile
+        >>> import xarray as xr
+        >>> data = xr.Dataset(
+        ...     coords={"index": [0, 1, 2]},
+        ...     data_vars={
+        ...         "x": ("index", [9, 8, 7]),
+        ...         "y": ("index", [6, 5, 4]),
+        ...         "z": ("index", [3, 2, 1]),
+        ...     },
+        ... )
+        >>> with Session() as ses:
+        ...     with ses.virtualfile_from_data(
+        ...         check_kind="vector", data=data
+        ...     ) as fin:
+        ...         # Send the output to a file so that we can read it
+        ...         with GMTTempFile() as fout:
+        ...             ses.call_module("info", f"{fin} ->{fout.name}")
+        ...             print(fout.read().strip())
+        ...
+        <vector memory>: N = 3 <7/9> <4/6> <1/3>
+        """
+        kind = data_kind(data, x, y, z)
+
+        if check_kind == "raster" and kind not in ("file", "grid"):
+            raise GMTInvalidInput(f"Unrecognized data type: {type(data)}")
+        if check_kind == "vector" and kind not in ("file", "matrix", "vectors"):
+            raise GMTInvalidInput(f"Unrecognized data type: {type(data)}")
+
+        # Decide which virtualfile_from_ function to use
+        _virtualfile_from = {
+            "file": dummy_context,
+            "grid": self.virtualfile_from_grid,
+            # Note: virtualfile_from_matrix is not used because a matrix can be
+            # converted to vectors instead, and using vectors allows for better
+            # handling of string type inputs (e.g. for datetime data types)
+            "matrix": self.virtualfile_from_vectors,
+            "vectors": self.virtualfile_from_vectors,
+        }[kind]
+
+        # Ensure the data is an iterable (Python list or tuple)
+        if kind in ("file", "grid"):
+            _data = (data,)
+        elif kind == "vectors":
+            _data = (x, y, z)
+        elif kind == "matrix":  # turn 2D arrays into list of vectors
+            try:
+                # pandas.DataFrame and xarray.Dataset types
+                _data = [array for _, array in data.items()]
+            except AttributeError:
+                # Python lists, tuples, and numpy ndarray types
+                _data = np.atleast_2d(np.asanyarray(data).T)
+
+        # Finally create the virtualfile from the data, to be passed into GMT
+        file_context = _virtualfile_from(*_data)
+
+        return file_context
+
     def extract_region(self):
         """
         Extract the WESN bounding box of the currently active figure.

pygmt/src/grdinfo.py

@@ -2,12 +2,9 @@
 grdinfo - Retrieve info about grid file.
 """
 from pygmt.clib import Session
-from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
     GMTTempFile,
     build_arg_string,
-    data_kind,
-    dummy_context,
     fmt_docstring,
     kwargs_to_strings,
     use_alias,
@@ -109,15 +106,9 @@ def grdinfo(grid, **kwargs):
     info : str
         A string with information about the grid.
     """
-    kind = data_kind(grid, None, None)
     with GMTTempFile() as outfile:
         with Session() as lib:
-            if kind == "file":
-                file_context = dummy_context(grid)
-            elif kind == "grid":
-                file_context = lib.virtualfile_from_grid(grid)
-            else:
-                raise GMTInvalidInput("Unrecognized data type: {}".format(type(grid)))
+            file_context = lib.virtualfile_from_data(check_kind="raster", data=grid)
             with file_context as infile:
                 arg_str = " ".join(
                     [infile, build_arg_string(kwargs), "->" + outfile.name]

pygmt/src/info.py

@@ -3,15 +3,7 @@
 """
 import numpy as np
 from pygmt.clib import Session
-from pygmt.exceptions import GMTInvalidInput
-from pygmt.helpers import (
-    GMTTempFile,
-    build_arg_string,
-    data_kind,
-    dummy_context,
-    fmt_docstring,
-    use_alias,
-)
+from pygmt.helpers import GMTTempFile, build_arg_string, fmt_docstring, use_alias
 
 
 @fmt_docstring
@@ -66,21 +58,8 @@ def info(table, **kwargs):
         - :class:`numpy.ndarray` if either of the above parameters are used.
         - str if none of the above parameters are used.
     """
-    kind = data_kind(table)
     with Session() as lib:
-        if kind == "file":
-            file_context = dummy_context(table)
-        elif kind == "matrix":
-            try:
-                # pandas.DataFrame and xarray.Dataset types
-                arrays = [array for _, array in table.items()]
-            except AttributeError:
-                # Python lists, tuples, and numpy ndarray types
-                arrays = np.atleast_2d(np.asanyarray(table).T)
-            file_context = lib.virtualfile_from_vectors(*arrays)
-        else:
-            raise GMTInvalidInput(f"Unrecognized data type: {type(table)}")
-
+        file_context = lib.virtualfile_from_data(data=table)
         with GMTTempFile() as tmpfile:
             with file_context as fname:
                 arg_str = " ".join(