Skip to content

Commit ea90891

Browse files
committed
Refactor the data_kind function and improve docstrings
1 parent d7560fa commit ea90891

File tree

1 file changed

+55
-40
lines changed

1 file changed

+55
-40
lines changed

pygmt/helpers/utils.py

Lines changed: 55 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -188,30 +188,36 @@ def _check_encoding(
188188

189189

190190
def data_kind(
191-
data: Any = None, required: bool = True
191+
data: Any, required: bool = True
192192
) -> Literal["arg", "file", "geojson", "grid", "image", "matrix", "vectors"]:
193193
"""
194194
Check the kind of data that is provided to a module.
195195
196-
The ``data`` argument can be in any type, but only following types are supported:
196+
Recognized data kinds are:
197197
198-
- a string or a :class:`pathlib.PurePath` object or a sequence of them, representing
199-
a file name or a list of file names
200-
- a 2-D or 3-D :class:`xarray.DataArray` object
201-
- a 2-D matrix
202-
- None, bool, int or float type representing an optional arguments
203-
- a geo-like Python object that implements ``__geo_interface__`` (e.g.,
204-
geopandas.GeoDataFrame or shapely.geometry)
198+
- ``"arg"``: bool, int or float, representing an optional argument, mainly used for
199+
dealing with optional virtual files
200+
- ``"file"``: a string or a :class:`pathlib.PurePath` object or a sequence of them,
201+
representing a file name or a list of file names
202+
- ``"geojson"``: a geo-like Python object that implements ``__geo_interface__``
203+
(e.g., geopandas.GeoDataFrame or shapely.geometry)
204+
- ``"grid"``: a :class:`xarray.DataArray` object with dimensions not equal to 3
205+
- ``"image"``: a :class:`xarray.DataArray` object with 3 dimensions
206+
- ``"matrix"``: a 2-D matrix or a sequence of sequences
207+
208+
In addition, the data can be given via a series of vectors (i.e., x/y/z). In this
209+
case, the ``data`` argument is ``None``. The data kind is ``"vectors"`` if the data
210+
is required. If the data is optional, the kind is ``"arg"``.
211+
212+
The function will fallback to ``"matrix"`` for any unrecognized data.
205213
206214
Parameters
207215
----------
208-
data : str, pathlib.PurePath, None, bool, xarray.DataArray or {table-like}
209-
Pass in either a file name or :class:`pathlib.Path` to an ASCII data
210-
table, an :class:`xarray.DataArray`, a 1-D/2-D
211-
{table-classes} or an option argument.
216+
data
217+
The data that is provided to a module.
212218
required
213-
Set to True when 'data' is required, or False when dealing with
214-
optional virtual files. [Default is True].
219+
If the data is required or not. Set to ``False`` when dealing with optional
220+
virtual files.
215221
216222
Returns
217223
-------
@@ -223,45 +229,54 @@ def data_kind(
223229
>>> import numpy as np
224230
>>> import xarray as xr
225231
>>> import pathlib
232+
>>> [data_kind(data=data) for data in (2, 2.0, True, False)]
233+
['arg', 'arg', 'arg', 'arg']
226234
>>> data_kind(data=None)
227235
'vectors'
228-
>>> data_kind(data=np.arange(10).reshape((5, 2)))
229-
'matrix'
236+
>>> data_kind(data=None, required=False)
237+
'arg'
230238
>>> data_kind(data="my-data-file.txt")
231239
'file'
232240
>>> data_kind(data=pathlib.Path("my-data-file.txt"))
233241
'file'
234-
>>> data_kind(data=None, required=False)
235-
'arg'
236-
>>> data_kind(data=2.0, required=False)
237-
'arg'
238-
>>> data_kind(data=True, required=False)
239-
'arg'
242+
>>> data_kind(data=["data1.txt", "data2.txt"])
243+
'file'
240244
>>> data_kind(data=xr.DataArray(np.random.rand(4, 3)))
241245
'grid'
242246
>>> data_kind(data=xr.DataArray(np.random.rand(3, 4, 5)))
243247
'image'
248+
>>> data_kind(data=np.arange(10).reshape((5, 2)))
249+
'matrix'
250+
>>> data_kind(data=[1, 2, 3])
251+
'matrix'
244252
"""
245-
kind: Literal["arg", "file", "geojson", "grid", "image", "matrix", "vectors"]
253+
# data is None, so data must be given via a series of vectors (i.e., x/y/z).
254+
# The only exception is when dealing with optional virtual files.
255+
if data is None:
256+
return "vectors" if required else "arg"
257+
258+
# A file or a list of files
246259
if isinstance(data, str | pathlib.PurePath) or (
247-
isinstance(data, list | tuple)
260+
is_nonstr_iter(data)
248261
and all(isinstance(_file, str | pathlib.PurePath) for _file in data)
249262
):
250-
# One or more files
251-
kind = "file"
252-
elif isinstance(data, bool | int | float) or (data is None and not required):
253-
kind = "arg"
254-
elif isinstance(data, xr.DataArray):
255-
kind = "image" if len(data.dims) == 3 else "grid"
256-
elif hasattr(data, "__geo_interface__"):
257-
# geo-like Python object that implements ``__geo_interface__``
258-
# (geopandas.GeoDataFrame or shapely.geometry)
259-
kind = "geojson"
260-
elif data is not None:
261-
kind = "matrix"
262-
else:
263-
kind = "vectors"
264-
return kind
263+
return "file"
264+
265+
# An option argument
266+
if isinstance(data, bool | int | float):
267+
return "arg"
268+
269+
# A xr.DataArray grid or image
270+
if isinstance(data, xr.DataArray):
271+
return "image" if len(data.dims) == 3 else "grid"
272+
273+
# Geo-like Python object that implements ``__geo_interface__`` (e.g.,
274+
# geopandas.GeoDataFrame or shapely.geometry)
275+
if hasattr(data, "__geo_interface__"):
276+
return "geojson"
277+
278+
# Fallback to "matrix" for anything else
279+
return "matrix"
265280

266281

267282
def non_ascii_to_octal(

0 commit comments

Comments
 (0)