Skip to content

Commit c024041

Browse files
committed
Refactor the data_kind function
1 parent f7f3d77 commit c024041

File tree

1 file changed

+39
-30
lines changed

1 file changed

+39
-30
lines changed

pygmt/helpers/utils.py

+39-30
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
import webbrowser
1111
from collections.abc import Iterable
1212

13+
import numpy as np
14+
import pandas as pd
1315
import xarray as xr
1416
from pygmt.exceptions import GMTInvalidInput
1517

@@ -21,18 +23,18 @@ def data_kind(data, x=None, y=None, z=None, required_z=False, optional_data=Fals
2123
Possible types:
2224
2325
* a file name provided as 'data'
24-
* an option argument provided as 'data'
25-
* a pathlib.Path provided as 'data'
26-
* an xarray.DataArray provided as 'data'
27-
* a matrix provided as 'data'
26+
* a pathlib.Path object provided as 'data'
27+
* an xarray.DataArray object provided as 'data'
28+
* a 2-D matrix provided as 'data'
2829
* 1-D arrays x and y (and z, optionally)
30+
* an option argument (None, bool, int, float or str type) provided as 'data'
2931
3032
Arguments should be ``None`` if not used. If doesn't fit any of these
3133
categories (or fits more than one), will raise an exception.
3234
3335
Parameters
3436
----------
35-
data : str or pathlib.Path or xarray.DataArray or {table-like} or None
37+
data : str or pathlib.Path or xarray.DataArray or {table-like} or None or bool
3638
Pass in either a file name or :class:`pathlib.Path` to an ASCII data
3739
table, an :class:`xarray.DataArray`, a 1-D/2-D
3840
{table-classes} or an option argument.
@@ -43,12 +45,13 @@ def data_kind(data, x=None, y=None, z=None, required_z=False, optional_data=Fals
4345
required_z : bool
4446
State whether the 'z' column is required.
4547
optional_data : bool
46-
State whether the 'data' is optional.
48+
State whether 'data' is optional (useful for dealing with optional
49+
virtual files).
4750
4851
Returns
4952
-------
5053
kind : str
51-
One of: ``'file'``, ``'grid'``, ``'geojson'``, ``'matrix'``, or
54+
One of ``'file'``, ``'grid'``, ``'geojson'``, ``'matrix'``, or
5255
``'vectors'``.
5356
5457
Examples
@@ -70,31 +73,37 @@ def data_kind(data, x=None, y=None, z=None, required_z=False, optional_data=Fals
7073
>>> data_kind(data=xr.DataArray(np.random.rand(4, 3)))
7174
'grid'
7275
"""
73-
if data is None and not optional_data and x is None and y is None:
74-
raise GMTInvalidInput("No input data provided.")
75-
if data is not None and (x is not None or y is not None or z is not None):
76-
raise GMTInvalidInput("Too much data. Use either data or x and y.")
77-
if data is None and not optional_data and (x is None or y is None):
76+
# validate the combinations of data/x/y/z
77+
if x is None and y is None: # both x and y are not given
78+
if data is None and not optional_data:
79+
raise GMTInvalidInput("No input data provides")
80+
elif x is None or y is None: # either x or y is not given
7881
raise GMTInvalidInput("Must provide both x and y.")
79-
if data is None and required_z and z is None:
80-
raise GMTInvalidInput("Must provide x, y, and z.")
81-
82-
if x is not None or y is not None or z is not None:
83-
kind = "vectors"
82+
else: # both x and y are given
83+
if data is not None:
84+
raise GMTInvalidInput("Too much data. Use either data or x and y.")
85+
if required_z and z is None:
86+
raise GMTInvalidInput("Must provide x, y, and z.")
87+
88+
# determine the data kind
89+
if isinstance(data, (bool, int, float, str, pathlib.Path)):
90+
# a null context manager will be created for "file-like" kind
91+
kind = "file"
92+
elif isinstance(data, xr.DataArray):
93+
kind = "grid"
94+
elif hasattr(data, "__geo_interface__"):
95+
# geo-like Python object that implements ``__geo_interface__``
96+
# (geopandas.GeoDataFrame or shapely.geometry)
97+
kind = "geojson"
98+
elif data is not None:
99+
if required_z and (
100+
getattr(data, "shape", (3, 3))[1] < 3 # np.array, pd.DataFrame
101+
or len(getattr(data, "data_vars", (0, 1, 2))) < 3 # xr.Dataset
102+
):
103+
raise GMTInvalidInput("data must provide x, y, and z columns.")
104+
kind = "matrix"
84105
else:
85-
if data is None or isinstance(data, (bool, int, float, str, pathlib.PurePath)):
86-
kind = "file"
87-
elif isinstance(data, xr.DataArray):
88-
kind = "grid"
89-
elif hasattr(data, "__geo_interface__"):
90-
kind = "geojson"
91-
else:
92-
if required_z and (
93-
getattr(data, "shape", (3, 3))[1] < 3 # np.array, pd.DataFrame
94-
or len(getattr(data, "data_vars", (0, 1, 2))) < 3 # xr.Dataset
95-
):
96-
raise GMTInvalidInput("data must provide x, y, and z columns.")
97-
kind = "matrix"
106+
kind = "vectors"
98107
return kind
99108

100109

0 commit comments

Comments
 (0)