Hypothesis strategies in xarray.testing.strategies

ci/requirements/doc.yml

@@ -10,6 +10,7 @@ dependencies:
   - cfgrib>=0.9
   - dask-core>=2.30
   - h5netcdf>=0.7.4
+  - hypothesis
   - ipykernel
   - ipython
   - iris>=2.3

doc/api.rst

@@ -1060,6 +1060,28 @@ Testing
    testing.assert_allclose
    testing.assert_chunks_equal
 
+Hypothesis Testing Strategies
+=============================
+
+.. currentmodule:: xarray
+
+.. warning::
+    These strategies should be considered highly experimental, and liable to change at any time.
+
+.. autosummary::
+   :toctree: generated/
+
+   testing.strategies.numeric_dtypes
+   testing.strategies.names
+   testing.strategies.dimension_names
+   testing.strategies.dimension_sizes
+   testing.strategies.attrs
+   testing.strategies.variables
+   testing.strategies.coordinate_variables
+   testing.strategies.dataarrays
+   testing.strategies.data_variables
+   testing.strategies.datasets
+
 Exceptions
 ==========
 

doc/user-guide/index.rst

@@ -25,4 +25,5 @@ examples that describe many common tasks that you can accomplish with xarray.
    dask
    plotting
    options
+   testing
    duckarrays

doc/user-guide/testing.rst

@@ -0,0 +1,272 @@
+.. _testing:
+
+Testing your code
+=================
+
+.. ipython:: python
+    :suppress:
+
+    import numpy as np
+    import pandas as pd
+    import xarray as xr
+
+    np.random.seed(123456)
+
+.. _hypothesis:
+
+Hypothesis testing
+------------------
+
+.. note::
+
+  Testing with hypothesis is a fairly advanced topic. Before reading this section it is recommended that you take a look
+  at our guide to xarray's :ref:`data structures`, are familiar with conventional unit testing in pytest, and have seen
+  the hypothesis library documentation.
+
+``Hypothesis`` is a powerful library for property-based testing.
+Instead of writing tests for one example at a time, it allows you to write tests parameterized by a source of many
+dynamically generated examples. For example you might have written a test which you wish to be parameterized by the set
+of all possible ``integers()``.
+
+Property-based testing is extremely powerful, because (unlike more conventional example-based testing) it can find bugs
+that you did not even think to look for!
+
+Strategies
+~~~~~~~~~~
+
+Each source of examples is called a "strategy", and xarray provides a range of custom strategies which produce xarray
+data structures containing arbitrary data. You can use these to efficiently test downstream code,
+quickly ensuring that your code can handle xarray objects of all possible structures and contents.
+
+These strategies are accessible in the :py:module::`xarray.testing.strategies` module, which provides
+
+.. currentmodule:: xarray
+
+.. autosummary::
+
+   testing.strategies.numeric_dtypes
+   testing.strategies.np_arrays
+   testing.strategies.names
+   testing.strategies.dimension_names
+   testing.strategies.dimension_sizes
+   testing.strategies.attrs
+   testing.strategies.variables
+   testing.strategies.coordinate_variables
+   testing.strategies.dataarrays
+   testing.strategies.data_variables
+   testing.strategies.datasets
+
+Generating Examples
+~~~~~~~~~~~~~~~~~~~
+
+To see an example of what each of these strategies might produce, you can call one followed by the ``.example()`` method,
+which is a general hypothesis method valid for all strategies.
+
+.. ipython:: python
+
+    import xarray.testing.strategies as xrst
+
+    xrst.dataarrays().example()
+    xrst.dataarrays().example()
+    xrst.dataarrays().example()
+
+You can see that calling ``.example()`` multiple times will generate different examples, giving you an idea of the wide
+range of data that the xarray strategies can generate.
+
+In your tests however you should not use ``.example()`` - instead you should parameterize your tests with the
+``hypothesis.given`` decorator:
+
+.. ipython:: python
+
+    from hypothesis import given
+
+.. ipython:: python
+
+    @given(xrst.dataarrays())
+    def test_function_that_acts_on_dataarrays(da):
+        assert func(da) == ...
+
+
+Chaining Strategies
+~~~~~~~~~~~~~~~~~~~
+
+Xarray's strategies can accept other strategies as arguments, allowing you to customise the contents of the generated
+examples.
+
+.. ipython:: python
+
+    # generate a DataArray with shape (3, 4), but all other details still arbitrary
+    xrst.dataarrays(
+        data=xrst.np_arrays(shape=(3, 4), dtype=np.dtype("int32"))
+    ).example()
+
+This also works with custom strategies, or strategies defined in other packages.
+For example you could create a ``chunks`` strategy to specify particular chunking patterns for a dask-backed array.
+
+.. warning::
+    When passing multiple different strategies to the same constructor the drawn examples must be mutually compatible.
+
+    In order to construct a valid xarray object to return, our strategies must check that the
+    variables / dimensions / coordinates are mutually compatible. If you pass multiple custom strategies to a strategy
+    constructor which are not compatible in all cases, an error will be raised, *even if they are still compatible in
+    other cases*. For example
+
+    .. code-block::
+
+        @given(st.data())
+        def test_something_else_inefficiently(data):
+            arrs = npst.arrays(dtype=numeric_dtypes)  # generates arrays of any shape
+            dims = xrst.dimension_names()  # generates lists of any number of dimensions
+
+            # Drawing examples from this strategy will raise a hypothesis.errors.InvalidArgument error.
+            var = data.draw(xrst.variables(data=arrs, dims=dims))
+
+            assert ...
+
+    Here we have passed custom strategies which won't often be compatible: only rarely will the array's ``ndims``
+    correspond to the number of dimensions drawn. We forbid arguments that are only *sometimes* compatible in order to
+    avoid extremely poor example generation performance (as generating invalid examples and rejecting them is
+    potentially unboundedly inefficient).
+
+
+Fixing Arguments
+~~~~~~~~~~~~~~~~
+
+If you want to fix one aspect of the data structure, whilst allowing variation in the generated examples
+over all other aspects, then use ``hypothesis.strategies.just()``.
+
+.. ipython:: python
+
+    import hypothesis.strategies as st
+
+    # Generates only dataarrays with dimensions ["x", "y"]
+    xrst.dataarrays(dims=st.just(["x", "y"])).example()
+
+(This is technically another example of chaining strategies - ``hypothesis.strategies.just`` is simply a special
+strategy that just contains a single example.)
+
+To fix the length of dimensions you can instead pass `dims` as a mapping of dimension names to lengths
+(i.e. following xarray objects' ``.sizes()`` property), e.g.
+
+.. ipython:: python
+
+    # Generates only dataarrays with dimensions ["x", "y"], of lengths 2 & 3 respectively
+    xrst.dataarrays(dims=st.just({"x": 2, "y": 3})).example()
+
+You can also use this to specify that you want examples which are missing some part of the data structure, for instance
+
+.. ipython:: python
+
+    # Generates only dataarrays with no coordinates
+    xrst.datasets(data_vars=st.just({})).example()
+
+Through a combination of chaining strategies and fixing arguments, you can specify quite complicated requirements on the
+objects your chained strategy will generate.
+
+.. ipython:: python
+
+    fixed_x_variable_y_maybe_z = st.fixed_dictionaries(
+        {"x": st.just(2), "y": st.integers(3, 4)}, optional={"z": st.just(2)}
+    )
+
+    fixed_x_variable_y_maybe_z.example()
+
+    special_dataarrays = xrst.dataarrays(dims=fixed_x_variable_y_maybe_z)
+
+    special_dataarrays.example()
+    special_dataarrays.example()
+
+Here we have used one of hypothesis' built-in strategies ``fixed_dictionaries`` to create a strategy which generates
+mappings of dimension names to lengths (i.e. the ``size`` of the xarray object we want).
+This particular strategy will always generate an ``x`` dimension of length 2, and a ``y`` dimension of
+length either 3 or 4, and will sometimes also generate a ``z`` dimension of length 2.
+By feeding this strategy for dictionaries into the `dims` argument of xarray's `dataarrays` strategy, we can generate
+arbitrary ``DataArray`` objects whose dimensions will always match these specifications.
+
+
+Creating Duck-type Arrays
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Xarray objects don't have to wrap numpy arrays, in fact they can wrap any array type which presents the same API as a
+numpy array (so-called "duck array wrapping", see :ref:`internals.duck_arrays`).
+
+Imagine we want to write a strategy which generates arbitrary `DataArray` objects, each of which wraps a
+``sparse.COO`` array instead of a ``numpy.ndarray``. How could we do that? There are two ways:
+
+1. Create a xarray object with numpy data and use ``.map()`` to convert the underlying array to a
+different type:
+
+.. ipython:: python
+    :okexcept:
+
+    import sparse
+    import hypothesis.extra.numpy as npst
+
+.. ipython:: python
+    :okexcept:
+
+    def convert_to_sparse(da):
+        if da.ndim == 0:
+            return da
+        else:
+            da.data = sparse.COO.from_numpy(da.values)
+            return da
+
+.. ipython:: python
+    :okexcept:
+
+    sparse_dataarrays = xrst.dataarrays().map(convert_to_sparse)
+
+    sparse_dataarrays.example()
+    sparse_dataarrays.example()
+
+2. Pass a strategy which generates the duck-typed arrays directly to the ``data`` argument of the xarray
+strategies:
+
+.. ipython:: python
+    :okexcept:
+
+    @st.composite
+    def sparse_arrays(draw) -> st.SearchStrategy[sparse._coo.core.COO]:
+        """Strategy which generates random sparse.COO arrays"""
+        shape = draw(npst.array_shapes())
+        density = draw(st.integers(min_value=0, max_value=1))
+        return sparse.random(shape, density=density)
+
+.. ipython:: python
+    :okexcept:
+
+    sparse_dataarrays = xrst.dataarrays(data=sparse_arrays())
+
+    sparse_dataarrays.example()
+    sparse_dataarrays.example()
+
+Either approach is fine, but one may be more convenient than the other depending on the type of the duck array which you
+want to wrap.
+
+Creating datasets can be a little more involved. Using method (1) is simple:
+
+.. ipython:: python
+    :okexcept:
+
+    def convert_ds_to_sparse(ds):
+        return ds.map(convert_to_sparse)
+
+.. ipython:: python
+    :okexcept:
+
+    sparse_datasets = xrst.datasets().map(convert_ds_to_sparse)
+
+    sparse_datasets.example()
+
+but building a dataset from scratch (i.e. method (2)) requires building the dataset object in such as way that all of
+the data variables have compatible dimensions. You can build up a dictionary of the form ``{var_name: data_variable}``
+yourself, or you can use the ``data_vars`` argument to the ``data_variables`` strategy (TODO):
+
+.. ipython:: python
+    :okexcept:
+
+    sparse_data_vars = xrst.data_variables(data=sparse_arrays())
+    sparse_datasets = xrst.datasets(data_vars=sparse_data_vars)
+
+    sparse_datasets.example()

doc/whats-new.rst

@@ -22,6 +22,10 @@ v2022.07.0 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Added a suite of hypothesis strategies for generating xarray objects containing arbitrary data, useful for testing.
+  Accessible under :py:func:`testing.strategies`, and documented in a new page on testing in the User Guide.
+  (:issue:`6911`, :pull:`6908`)
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/testing/__init__.py

@@ -0,0 +1,23 @@
+from .testing import (  # noqa: F401
+    _assert_dataarray_invariants,
+    _assert_dataset_invariants,
+    _assert_indexes_invariants_checks,
+    _assert_internal_invariants,
+    _assert_variable_invariants,
+    _data_allclose_or_equiv,
+    assert_allclose,
+    assert_chunks_equal,
+    assert_duckarray_allclose,
+    assert_duckarray_equal,
+    assert_equal,
+    assert_identical,
+)
+
+__all__ = [
+    "assert_allclose",
+    "assert_chunks_equal",
+    "assert_duckarray_equal",
+    "assert_duckarray_allclose",
+    "assert_equal",
+    "assert_identical",
+]