Repository: fatiando/verde
Pull request #295 (2020-10-27 19:19:43)

Refactor get_data_names and check functions
Refactor `get_data_names` as a method of `BaseGridder`: simplify
`get_data_names` logic, make `get_data_names` to use `check_data_names` to
convert single string to a tuple and check that `data_names` has the same
number of elements as data. Add `data_names_defaults` as a class attribute of
`BaseGridder`. Move `check_extra_coords_names` to `verde/base/utils.py` along
with the other check functions. Make `check_data_names` to raise error if
`data_names` is `None`. But `BaseGridder.get_data_names` will return default
values if `data_names` is `None`. Rename test functions for `make_xarray_grid`.



**Reminders**:

- [ ] Run `make format` and `make check` to make sure the code follows the style guide.
- [ ] Add tests for new features or tests that would have caught the bug that you're fixing.
- [ ] Add new public functions/methods/classes to `doc/api/index.rst` and the base `__init__.py` file for the package.
- [ ] Write detailed docstrings for all functions/classes/methods. It often helps to design better code if you write the docstrings first.
- [ ] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.
- [ ] Add your full name, affiliation, and [ORCID](https://orcid.org) (optional) to the `AUTHORS.md` file (if you haven't already) in case you'd like to be listed as an author on the [Zenodo](https://zenodo.org/communities/fatiando) archive of the next release.



Patch:
diff --git a/verde/base/base_classes.py b/verde/base/base_classes.py
index 18f599b..0061025 100644
--- a/verde/base/base_classes.py
+++ b/verde/base/base_classes.py
@@ -209,6 +209,14 @@ class BaseGridder(BaseEstimator):
     # using this name as a basis.
     extra_coords_name = "extra_coord"
 
+    # Define default values for data_names depending on the number of data
+    # arrays returned by predict method.
+    data_names_defaults = [
+        ("scalars",),
+        ("east_component", "north_component"),
+        ("east_component", "north_component", "vertical_component"),
+    ]
+
     def predict(self, coordinates):
         """
         Predict data on the given coordinate values. NOT IMPLEMENTED.
@@ -416,8 +424,7 @@ class BaseGridder(BaseEstimator):
                 self.predict(project_coordinates(coordinates, projection))
             )
         # Get names for data and any extra coordinates
-        data_names = check_data_names(data_names)
-        data_names = get_data_names(data, data_names)
+        data_names = self._get_data_names(data, data_names)
         extra_coords_names = self._get_extra_coords_names(coordinates)
         # Create xarray.Dataset
         dataset = make_xarray_grid(
@@ -506,8 +513,7 @@ class BaseGridder(BaseEstimator):
             data = check_data(
                 self.predict(project_coordinates(coordinates, projection))
             )
-        data_names = check_data_names(data_names)
-        data_names = get_data_names(data, data_names)
+        data_names = self._get_data_names(data, data_names)
         columns = [(dims[0], coordinates[1]), (dims[1], coordinates[0])]
         extra_coords_names = self._get_extra_coords_names(coordinates)
         columns.extend(zip(extra_coords_names, coordinates[2:]))
@@ -618,8 +624,7 @@ class BaseGridder(BaseEstimator):
         # profile but Cartesian distances.
         if projection is not None:
             coordinates = project_coordinates(coordinates, projection, inverse=True)
-        data_names = check_data_names(data_names)
-        data_names = get_data_names(data, data_names)
+        data_names = self._get_data_names(data, data_names)
         columns = [
             (dims[0], coordinates[1]),
             (dims[1], coordinates[0]),
@@ -665,6 +670,46 @@ class BaseGridder(BaseEstimator):
             names.append(name)
         return names
 
+    def _get_data_names(self, data, data_names):
+        """
+        Get default names for data fields if none are given based on the data.
+
+        Examples
+        --------
+
+        >>> import numpy as np
+        >>> east, north, up = [np.arange(10)]*3
+        >>> gridder = BaseGridder()
+        >>> gridder._get_data_names((east,), data_names=None)
+        ('scalars',)
+        >>> gridder._get_data_names((east, north), data_names=None)
+        ('east_component', 'north_component')
+        >>> gridder._get_data_names((east, north, up), data_names=None)
+        ('east_component', 'north_component', 'vertical_component')
+        >>> gridder._get_data_names((east,), data_names="john")
+        ('john',)
+        >>> gridder._get_data_names((east,), data_names=("paul",))
+        ('paul',)
+        >>> gridder._get_data_names(
+        ...     (up, north), data_names=('ringo', 'george')
+        ... )
+        ('ringo', 'george')
+        >>> gridder._get_data_names((north,), data_names=["brian"])
+        ['brian']
+
+        """
+        # Return the defaults data_names for the class
+        if data_names is None:
+            if len(data) > len(self.data_names_defaults):
+                raise ValueError(
+                    "Default data names only available for up to 3 components. "
+                    + "Must provide custom names through the 'data_names' argument."
+                )
+            return self.data_names_defaults[len(data) - 1]
+        # Return the passed data_names if valid
+        data_names = check_data_names(data, data_names)
+        return data_names
+
 
 def project_coordinates(coordinates, projection, **kwargs):
     """
@@ -702,50 +747,6 @@ def project_coordinates(coordinates, projection, **kwargs):
     return proj_coordinates
 
 
-def get_data_names(data, data_names):
-    """
-    Get default names for data fields if none are given based on the data.
-
-    Examples
-    --------
-
-    >>> import numpy as np
-    >>> east, north, up = [np.arange(10)]*3
-    >>> get_data_names((east,), data_names=None)
-    ('scalars',)
-    >>> get_data_names((east, north), data_names=None)
-    ('east_component', 'north_component')
-    >>> get_data_names((east, north, up), data_names=None)
-    ('east_component', 'north_component', 'vertical_component')
-    >>> get_data_names((up, north), data_names=('ringo', 'george'))
-    ('ringo', 'george')
-
-    """
-    if data_names is not None:
-        if len(data) != len(data_names):
-            raise ValueError(
-                "Data has {} components but only {} names provided: {}".format(
-                    len(data), len(data_names), str(data_names)
-                )
-            )
-        return data_names
-    data_types = [
-        ("scalars",),
-        ("east_component", "north_component"),
-        ("east_component", "north_component", "vertical_component"),
-    ]
-    if len(data) > len(data_types):
-        raise ValueError(
-            " ".join(
-                [
-                    "Default data names only available for up to 3 components.",
-                    "Must provide custom names through the 'data_names' argument.",
-                ]
-            )
-        )
-    return data_types[len(data) - 1]
-
-
 def get_instance_region(instance, region):
     """
     Get the region attribute stored in instance if one is not provided.
diff --git a/verde/base/utils.py b/verde/base/utils.py
index 8cee240..3f40f1e 100644
--- a/verde/base/utils.py
+++ b/verde/base/utils.py
@@ -104,27 +104,39 @@ def check_data(data):
     return data
 
 
-def check_data_names(data_names):
+def check_data_names(data, data_names):
     """
-    Check the *data_names* argument and make sure it's a tuple.
-    If ``data_names`` is a single string, return it as a tuple with a single
-    element.
+    Check *data_names* against *data*.
 
-    This is the default form accepted by gridders and functions that require
-    the ``data_names`` argument.
+    Also, convert ``data_names`` to a tuple if it's a single string.
 
     Examples
     --------
 
-    >>> check_data_names("dummy")
+    >>> import numpy as np
+    >>> east, north, scalar = [np.array(10)]*3
+    >>> check_data_names((scalar,), "dummy")
     ('dummy',)
-    >>> check_data_names(("component_x", "component_y"))
-    ('component_x', 'component_y')
-    >>> check_data_names(["dummy"])
+    >>> check_data_names((scalar,), ("dummy",))
+    ('dummy',)
+    >>> check_data_names((scalar,), ["dummy"])
     ['dummy']
+    >>> check_data_names((east, north), ("component_x", "component_y"))
+    ('component_x', 'component_y')
     """
+    # Convert single string to tuple
     if isinstance(data_names, str):
         data_names = (data_names,)
+    # Raise error if data_names is None
+    if data_names is None:
+        raise ValueError("Invalid data_names equal to None.")
+    # Raise error if data and data_names don't have the same number of elements
+    if len(data) != len(data_names):
+        raise ValueError(
+            "Data has {} components but only {} names provided: {}".format(
+                len(data), len(data_names), str(data_names)
+            )
+        )
     return data_names
 
 
@@ -143,6 +155,46 @@ def check_coordinates(coordinates):
     return coordinates
 
 
+def check_extra_coords_names(coordinates, extra_coords_names):
+    """
+    Check extra_coords_names against coordiantes.
+
+    Also, convert ``extra_coords_names`` to a tuple if it's a single string.
+    Assume that there are extra coordinates on the ``coordinates`` tuple.
+
+    Examples
+    --------
+
+    >>> import numpy as np
+    >>> coordinates = [np.array(10)]*3
+    >>> check_extra_coords_names(coordinates, "upward")
+    ('upward',)
+    >>> check_extra_coords_names(coordinates, ("upward",))
+    ('upward',)
+    >>> coordinates = [np.array(10)]*4
+    >>> check_extra_coords_names(coordinates, ("upward", "time"))
+    ('upward', 'time')
+    """
+    # Convert single string to a tuple
+    if isinstance(extra_coords_names, str):
+        extra_coords_names = (extra_coords_names,)
+    # Check if it's not None
+    if extra_coords_names is None:
+        raise ValueError(
+            "Invalid extra_coords_names equal to None. "
+            + "When passing one or more extra coordinate, "
+            + "extra_coords_names cannot be None."
+        )
+    # Check if there are the same number of extra_coords than extra_coords_name
+    if len(coordinates[2:]) != len(extra_coords_names):
+        raise ValueError(
+            "Invalid extra_coords_names '{}'. ".format(extra_coords_names)
+            + "Number of extra coordinates names must match the number of "
+            + "additional coordinates ('{}').".format(len(coordinates[2:]))
+        )
+    return extra_coords_names
+
+
 def check_fit_input(coordinates, data, weights, unpack=True):
     """
     Validate the inputs to the fit method of gridders.
diff --git a/verde/utils.py b/verde/utils.py
index 3c40859..6f18740 100644
--- a/verde/utils.py
+++ b/verde/utils.py
@@ -19,8 +19,13 @@ try:
 except ImportError:
     numba = None
 
-from .base.utils import check_data, check_data_names, n_1d_arrays
-from .coordinates import check_coordinates
+from .base.utils import (
+    check_coordinates,
+    check_extra_coords_names,
+    check_data,
+    check_data_names,
+    n_1d_arrays,
+)
 
 
 def dispatch(function, delayed=False, client=None):
@@ -244,9 +249,9 @@ def make_xarray_grid(
         Array or tuple of arrays with data values on each point in the grid.
         Each array must contain values for a dimension in the same order as
         the coordinates. All arrays need to have the same *shape*.
-    data_names : str, list or None
+    data_names : str or list
         The name(s) of the data variables in the output grid.
-    dims : list or None
+    dims : list
         The names of the northing and easting data dimensions, respectively,
         in the output grid. Must be defined in the following order: northing
         dimension, easting dimension.
@@ -308,53 +313,20 @@ def make_xarray_grid(
     """
     coordinates = check_coordinates(coordinates)
     data = check_data(data)
-    data_names = check_data_names(data_names)
+    data_names = check_data_names(data, data_names)
     # dims is like shape with order (rows, cols) for the array
     # so the first element is northing and second is easting
     coords = {dims[1]: coordinates[0][0, :], dims[0]: coordinates[1][:, 0]}
     # Extra coordinates are handled like 2D data arrays with
     # the same dims and the data.
     if coordinates[2:]:
-        extra_coords_names = _check_extra_coords_names(coordinates, extra_coords_names)
+        extra_coords_names = check_extra_coords_names(coordinates, extra_coords_names)
         for name, extra_coord in zip(extra_coords_names, coordinates[2:]):
             coords[name] = (dims, extra_coord)
-    # Generate object
-    if len(data) != len(data_names):
-        raise ValueError(
-            "Invalid data_names '{}'. ".format(data_names)
-            + "Number of data names must match the number of "
-            + "data arrays ('{}').".format(len(data))
-        )
     data_vars = {name: (dims, value) for name, value in zip(data_names, data)}
     return xr.Dataset(data_vars, coords)
 
 
-def _check_extra_coords_names(coordinates, extra_coords_names):
-    """
-    Sanity checks for extra_coords_names against coordinates
-
-    Assuming that there are extra coordinates on the ``coordinates`` tuple.
-    """
-    # Convert to tuple if it's a str
-    if isinstance(extra_coords_names, str):
-        extra_coords_names = (extra_coords_names,)
-    # Check if it's not None
-    if extra_coords_names is None:
-        raise ValueError(
-            "Invalid extra_coords_names equal to None. "
-            + "When passing one or more extra coordinate, "
-            + "extra_coords_names cannot be None."
-        )
-    # Check if there are the same number of extra_coords than extra_coords_name
-    if len(coordinates[2:]) != len(extra_coords_names):
-        raise ValueError(
-            "Invalid extra_coords_names '{}'. ".format(extra_coords_names)
-            + "Number of extra coordinates names must match the number of "
-            + "additional coordinates ('{}').".format(len(coordinates[2:]))
-        )
-    return extra_coords_names
-
-
 def grid_to_table(grid):
     """
     Convert a grid to a table with the values and coordinates of each point.


Test patch:
diff --git a/verde/tests/test_base.py b/verde/tests/test_base.py
index 7a0ee1e..5baa58b 100644
--- a/verde/tests/test_base.py
+++ b/verde/tests/test_base.py
@@ -11,7 +11,6 @@ from ..base.utils import check_fit_input, check_coordinates
 from ..base.base_classes import (
     BaseGridder,
     BaseBlockCrossValidator,
-    get_data_names,
     get_instance_region,
 )
 from ..coordinates import grid_coordinates, scatter_points
@@ -46,28 +45,30 @@ def test_get_data_names():
     data2 = tuple([np.arange(10)] * 2)
     data3 = tuple([np.arange(10)] * 3)
     # Test the default names
-    assert get_data_names(data1, data_names=None) == ("scalars",)
-    assert get_data_names(data2, data_names=None) == (
+    gridder = BaseGridder()
+    assert gridder._get_data_names(data1, data_names=None) == ("scalars",)
+    assert gridder._get_data_names(data2, data_names=None) == (
         "east_component",
         "north_component",
     )
-    assert get_data_names(data3, data_names=None) == (
+    assert gridder._get_data_names(data3, data_names=None) == (
         "east_component",
         "north_component",
         "vertical_component",
     )
     # Test custom names
-    assert get_data_names(data1, data_names=("a",)) == ("a",)
-    assert get_data_names(data2, data_names=("a", "b")) == ("a", "b")
-    assert get_data_names(data3, data_names=("a", "b", "c")) == ("a", "b", "c")
+    assert gridder._get_data_names(data1, data_names=("a",)) == ("a",)
+    assert gridder._get_data_names(data2, data_names=("a", "b")) == ("a", "b")
+    assert gridder._get_data_names(data3, data_names=("a", "b", "c")) == ("a", "b", "c")
 
 
 def test_get_data_names_fails():
     "Check if fails for invalid data types"
+    gridder = BaseGridder()
     with pytest.raises(ValueError):
-        get_data_names(tuple([np.arange(5)] * 4), data_names=None)
+        gridder._get_data_names(tuple([np.arange(5)] * 4), data_names=None)
     with pytest.raises(ValueError):
-        get_data_names(tuple([np.arange(5)] * 2), data_names=("meh",))
+        gridder._get_data_names(tuple([np.arange(5)] * 2), data_names=("meh",))
 
 
 def test_get_instance_region():
diff --git a/verde/tests/test_utils.py b/verde/tests/test_utils.py
index 3553c33..f1ea645 100644
--- a/verde/tests/test_utils.py
+++ b/verde/tests/test_utils.py
@@ -98,7 +98,7 @@ def test_partition_by_sum_fails_no_partitions():
     assert "Could not find partition points" in str(error)
 
 
-def test_build_grid():
+def test_make_xarray_grid():
     """
     Check if xarray.Dataset is correctly created
     """
@@ -121,7 +121,7 @@ def test_build_grid():
     assert grid.dummy.shape == (5, 6)
 
 
-def test_build_grid_multiple_data():
+def test_make_xarray_grid_multiple_data():
     """
     Check if xarray.Dataset with multiple data is correctly created
     """
@@ -138,7 +138,7 @@ def test_build_grid_multiple_data():
         assert dataset["data_{}".format(i)].shape == (5, 6)
 
 
-def test_build_grid_extra_coords():
+def test_make_xarray_grid_extra_coords():
     """
     Check if xarray.Dataset with extra coords is correctly created
     """
@@ -163,7 +163,7 @@ def test_build_grid_extra_coords():
     assert dataset.time.shape == (5, 6)
 
 
-def test_build_grid_invalid_names():
+def test_make_xarray_grid_invalid_names():
     """
     Check if errors are raise after invalid data names
     """
@@ -174,13 +174,16 @@ def test_build_grid_invalid_names():
     data = np.ones_like(coordinates[0])
     with pytest.raises(ValueError):
         make_xarray_grid(coordinates, data, data_names=["bla_1", "bla_2"])
+    # data_names equal to None
+    with pytest.raises(ValueError):
+        make_xarray_grid(coordinates, data, data_names=None)
     # Multiple data, single data_name
     data = tuple(i * np.ones_like(coordinates[0]) for i in (1, 2))
     with pytest.raises(ValueError):
         make_xarray_grid(coordinates, data, data_names="blabla")
 
 
-def test_build_grid_invalid_extra_coords():
+def test_make_xarray_grid_invalid_extra_coords():
     """
     Check if errors are raise after invalid extra coords
     """