Merge branch 'main' into coarsen_reshape

dcherian · dcherian · commit 56d8c368591e · 2021-06-23T10:18:09.000-06:00
* main: Improve error message for guess engine (pydata#5455) Refactor dataset groupby tests (pydata#5506) DOC: zarr note on encoding (pydata#5427) Allow plotting categorical data (pydata#5464)
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -46,6 +46,8 @@ New Features
   By `Thomas Hirtz <https://github.com/thomashirtz>`_.
 - allow passing a function to ``combine_attrs`` (:pull:`4896`).
   By `Justus Magin <https://github.com/keewis>`_.
+- Allow plotting categorical data (:pull:`5464`).
+  By `Jimmy Westling <https://github.com/illviljan>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
diff --git a/xarray/backends/cfgrib_.py b/xarray/backends/cfgrib_.py
@@ -94,6 +94,8 @@ def get_encoding(self):
 
 
 class CfgribfBackendEntrypoint(BackendEntrypoint):
+    available = has_cfgrib
+
     def guess_can_open(self, filename_or_obj):
         try:
             _, ext = os.path.splitext(filename_or_obj)
@@ -147,5 +149,4 @@ def open_dataset(
         return ds
 
 
-if has_cfgrib:
-    BACKEND_ENTRYPOINTS["cfgrib"] = CfgribfBackendEntrypoint
+BACKEND_ENTRYPOINTS["cfgrib"] = CfgribfBackendEntrypoint
diff --git a/xarray/backends/h5netcdf_.py b/xarray/backends/h5netcdf_.py
@@ -337,6 +337,8 @@ def close(self, **kwargs):
 
 
 class H5netcdfBackendEntrypoint(BackendEntrypoint):
+    available = has_h5netcdf
+
     def guess_can_open(self, filename_or_obj):
         magic_number = try_read_magic_number_from_file_or_path(filename_or_obj)
         if magic_number is not None:
@@ -394,5 +396,4 @@ def open_dataset(
         return ds
 
 
-if has_h5netcdf:
-    BACKEND_ENTRYPOINTS["h5netcdf"] = H5netcdfBackendEntrypoint
+BACKEND_ENTRYPOINTS["h5netcdf"] = H5netcdfBackendEntrypoint
diff --git a/xarray/backends/netCDF4_.py b/xarray/backends/netCDF4_.py
@@ -512,6 +512,8 @@ def close(self, **kwargs):
 
 
 class NetCDF4BackendEntrypoint(BackendEntrypoint):
+    available = has_netcdf4
+
     def guess_can_open(self, filename_or_obj):
         if isinstance(filename_or_obj, str) and is_remote_uri(filename_or_obj):
             return True
@@ -573,5 +575,4 @@ def open_dataset(
         return ds
 
 
-if has_netcdf4:
-    BACKEND_ENTRYPOINTS["netcdf4"] = NetCDF4BackendEntrypoint
+BACKEND_ENTRYPOINTS["netcdf4"] = NetCDF4BackendEntrypoint
diff --git a/xarray/backends/plugins.py b/xarray/backends/plugins.py
@@ -81,7 +81,10 @@ def sort_backends(backend_entrypoints):
 
 
 def build_engines(pkg_entrypoints):
-    backend_entrypoints = BACKEND_ENTRYPOINTS.copy()
+    backend_entrypoints = {}
+    for backend_name, backend in BACKEND_ENTRYPOINTS.items():
+        if backend.available:
+            backend_entrypoints[backend_name] = backend
     pkg_entrypoints = remove_duplicates(pkg_entrypoints)
     external_backend_entrypoints = backends_dict_from_pkg(pkg_entrypoints)
     backend_entrypoints.update(external_backend_entrypoints)
@@ -101,30 +104,49 @@ def guess_engine(store_spec):
 
     for engine, backend in engines.items():
         try:
-            if backend.guess_can_open and backend.guess_can_open(store_spec):
+            if backend.guess_can_open(store_spec):
                 return engine
         except Exception:
             warnings.warn(f"{engine!r} fails while guessing", RuntimeWarning)
 
-    installed = [k for k in engines if k != "store"]
-    if installed:
-        raise ValueError(
-            "did not find a match in any of xarray's currently installed IO "
-            f"backends {installed}. Consider explicitly selecting one of the "
-            "installed backends via the ``engine`` parameter to "
-            "xarray.open_dataset(), or installing additional IO dependencies:\n"
-            "http://xarray.pydata.org/en/stable/getting-started-guide/installing.html\n"
-            "http://xarray.pydata.org/en/stable/user-guide/io.html"
-        )
+    compatible_engines = []
+    for engine, backend_cls in BACKEND_ENTRYPOINTS.items():
+        try:
+            backend = backend_cls()
+            if backend.guess_can_open(store_spec):
+                compatible_engines.append(engine)
+        except Exception:
+            warnings.warn(f"{engine!r} fails while guessing", RuntimeWarning)
+
+    installed_engines = [k for k in engines if k != "store"]
+    if not compatible_engines:
+        if installed_engines:
+            error_msg = (
+                "did not find a match in any of xarray's currently installed IO "
+                f"backends {installed_engines}. Consider explicitly selecting one of the "
+                "installed engines via the ``engine`` parameter, or installing "
+                "additional IO dependencies, see:\n"
+                "http://xarray.pydata.org/en/stable/getting-started-guide/installing.html\n"
+                "http://xarray.pydata.org/en/stable/user-guide/io.html"
+            )
+        else:
+            error_msg = (
+                "xarray is unable to open this file because it has no currently "
+                "installed IO backends. Xarray's read/write support requires "
+                "installing optional IO dependencies, see:\n"
+                "http://xarray.pydata.org/en/stable/getting-started-guide/installing.html\n"
+                "http://xarray.pydata.org/en/stable/user-guide/io"
+            )
     else:
-        raise ValueError(
-            "xarray is unable to open this file because it has no currently "
-            "installed IO backends. Xarray's read/write support requires "
-            "installing optional dependencies:\n"
-            "http://xarray.pydata.org/en/stable/getting-started-guide/installing.html\n"
-            "http://xarray.pydata.org/en/stable/user-guide/io.html"
+        error_msg = (
+            "found the following matches with the input file in xarray's IO "
+            f"backends: {compatible_engines}. But their dependencies may not be installed, see:\n"
+            "http://xarray.pydata.org/en/stable/user-guide/io.html \n"
+            "http://xarray.pydata.org/en/stable/getting-started-guide/installing.html"
         )
 
+    raise ValueError(error_msg)
+
 
 def get_backend(engine):
     """Select open_dataset method based on current engine."""
diff --git a/xarray/backends/pseudonetcdf_.py b/xarray/backends/pseudonetcdf_.py
@@ -102,6 +102,7 @@ def close(self):
 
 
 class PseudoNetCDFBackendEntrypoint(BackendEntrypoint):
+    available = has_pseudonetcdf
 
     # *args and **kwargs are not allowed in open_backend_dataset_ kwargs,
     # unless the open_dataset_parameters are explicity defined like this:
@@ -153,5 +154,4 @@ def open_dataset(
         return ds
 
 
-if has_pseudonetcdf:
-    BACKEND_ENTRYPOINTS["pseudonetcdf"] = PseudoNetCDFBackendEntrypoint
+BACKEND_ENTRYPOINTS["pseudonetcdf"] = PseudoNetCDFBackendEntrypoint
diff --git a/xarray/backends/pydap_.py b/xarray/backends/pydap_.py
@@ -110,6 +110,8 @@ def get_dimensions(self):
 
 
 class PydapBackendEntrypoint(BackendEntrypoint):
+    available = has_pydap
+
     def guess_can_open(self, filename_or_obj):
         return isinstance(filename_or_obj, str) and is_remote_uri(filename_or_obj)
 
@@ -154,5 +156,4 @@ def open_dataset(
             return ds
 
 
-if has_pydap:
-    BACKEND_ENTRYPOINTS["pydap"] = PydapBackendEntrypoint
+BACKEND_ENTRYPOINTS["pydap"] = PydapBackendEntrypoint
diff --git a/xarray/backends/pynio_.py b/xarray/backends/pynio_.py
@@ -99,6 +99,8 @@ def close(self):
 
 
 class PynioBackendEntrypoint(BackendEntrypoint):
+    available = has_pynio
+
     def open_dataset(
         self,
         filename_or_obj,
@@ -112,13 +114,13 @@ def open_dataset(
         mode="r",
         lock=None,
     ):
+        filename_or_obj = _normalize_path(filename_or_obj)
         store = NioDataStore(
             filename_or_obj,
             mode=mode,
             lock=lock,
         )
 
-        filename_or_obj = _normalize_path(filename_or_obj)
         store_entrypoint = StoreBackendEntrypoint()
         with close_on_error(store):
             ds = store_entrypoint.open_dataset(
@@ -134,5 +136,4 @@ def open_dataset(
         return ds
 
 
-if has_pynio:
-    BACKEND_ENTRYPOINTS["pynio"] = PynioBackendEntrypoint
+BACKEND_ENTRYPOINTS["pynio"] = PynioBackendEntrypoint
diff --git a/xarray/backends/scipy_.py b/xarray/backends/scipy_.py
@@ -238,6 +238,8 @@ def close(self):
 
 
 class ScipyBackendEntrypoint(BackendEntrypoint):
+    available = has_scipy
+
     def guess_can_open(self, filename_or_obj):
 
         magic_number = try_read_magic_number_from_file_or_path(filename_or_obj)
@@ -290,5 +292,4 @@ def open_dataset(
         return ds
 
 
-if has_scipy:
-    BACKEND_ENTRYPOINTS["scipy"] = ScipyBackendEntrypoint
+BACKEND_ENTRYPOINTS["scipy"] = ScipyBackendEntrypoint
diff --git a/xarray/backends/store.py b/xarray/backends/store.py
@@ -4,6 +4,8 @@
 
 
 class StoreBackendEntrypoint(BackendEntrypoint):
+    available = True
+
     def guess_can_open(self, filename_or_obj):
         return isinstance(filename_or_obj, AbstractDataStore)
 
diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
@@ -785,6 +785,15 @@ def open_zarr(
 
 
 class ZarrBackendEntrypoint(BackendEntrypoint):
+    available = has_zarr
+
+    def guess_can_open(self, filename_or_obj):
+        try:
+            _, ext = os.path.splitext(filename_or_obj)
+        except TypeError:
+            return False
+        return ext in {".zarr"}
+
     def open_dataset(
         self,
         filename_or_obj,
@@ -840,5 +849,4 @@ def open_dataset(
         return ds
 
 
-if has_zarr:
-    BACKEND_ENTRYPOINTS["zarr"] = ZarrBackendEntrypoint
+BACKEND_ENTRYPOINTS["zarr"] = ZarrBackendEntrypoint
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -2001,6 +2001,10 @@ def to_zarr(
             If not other chunks are found, Zarr uses its own heuristics to
             choose automatic chunk sizes.
 
+        encoding:
+            The encoding attribute (if exists) of the DataArray(s) will be
+            used. Override any existing encodings by providing the ``encoding`` kwarg.
+
         See Also
         --------
         :ref:`io.zarr`
diff --git a/xarray/plot/plot.py b/xarray/plot/plot.py
@@ -925,18 +925,26 @@ def imshow(x, y, z, ax, **kwargs):
             "imshow requires 1D coordinates, try using pcolormesh or contour(f)"
         )
 
-    # Centering the pixels- Assumes uniform spacing
-    try:
-        xstep = (x[1] - x[0]) / 2.0
-    except IndexError:
-        # Arbitrary default value, similar to matplotlib behaviour
-        xstep = 0.1
-    try:
-        ystep = (y[1] - y[0]) / 2.0
-    except IndexError:
-        ystep = 0.1
-    left, right = x[0] - xstep, x[-1] + xstep
-    bottom, top = y[-1] + ystep, y[0] - ystep
+    def _center_pixels(x):
+        """Center the pixels on the coordinates."""
+        if np.issubdtype(x.dtype, str):
+            # When using strings as inputs imshow converts it to
+            # integers. Choose extent values which puts the indices in
+            # in the center of the pixels:
+            return 0 - 0.5, len(x) - 0.5
+
+        try:
+            # Center the pixels assuming uniform spacing:
+            xstep = 0.5 * (x[1] - x[0])
+        except IndexError:
+            # Arbitrary default value, similar to matplotlib behaviour:
+            xstep = 0.1
+
+        return x[0] - xstep, x[-1] + xstep
+
+    # Center the pixels:
+    left, right = _center_pixels(x)
+    top, bottom = _center_pixels(y)
 
     defaults = {"origin": "upper", "interpolation": "nearest"}
 
@@ -967,6 +975,13 @@ def imshow(x, y, z, ax, **kwargs):
 
     primitive = ax.imshow(z, **defaults)
 
+    # If x or y are strings the ticklabels have been replaced with
+    # integer indices. Replace them back to strings:
+    for axis, v in [("x", x), ("y", y)]:
+        if np.issubdtype(v.dtype, str):
+            getattr(ax, f"set_{axis}ticks")(np.arange(len(v)))
+            getattr(ax, f"set_{axis}ticklabels")(v)
+
     return primitive
 
 
@@ -1011,9 +1026,13 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
         else:
             infer_intervals = True
 
-    if infer_intervals and (
-        (np.shape(x)[0] == np.shape(z)[1])
-        or ((x.ndim > 1) and (np.shape(x)[1] == np.shape(z)[1]))
+    if (
+        infer_intervals
+        and not np.issubdtype(x.dtype, str)
+        and (
+            (np.shape(x)[0] == np.shape(z)[1])
+            or ((x.ndim > 1) and (np.shape(x)[1] == np.shape(z)[1]))
+        )
     ):
         if len(x.shape) == 1:
             x = _infer_interval_breaks(x, check_monotonic=True)
@@ -1022,7 +1041,11 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
             x = _infer_interval_breaks(x, axis=1)
             x = _infer_interval_breaks(x, axis=0)
 
-    if infer_intervals and (np.shape(y)[0] == np.shape(z)[0]):
+    if (
+        infer_intervals
+        and not np.issubdtype(y.dtype, str)
+        and (np.shape(y)[0] == np.shape(z)[0])
+    ):
         if len(y.shape) == 1:
             y = _infer_interval_breaks(y, check_monotonic=True)
         else:
diff --git a/xarray/plot/utils.py b/xarray/plot/utils.py
@@ -604,7 +604,14 @@ def _ensure_plottable(*args):
     Raise exception if there is anything in args that can't be plotted on an
     axis by matplotlib.
     """
-    numpy_types = [np.floating, np.integer, np.timedelta64, np.datetime64, np.bool_]
+    numpy_types = [
+        np.floating,
+        np.integer,
+        np.timedelta64,
+        np.datetime64,
+        np.bool_,
+        np.str_,
+    ]
     other_types = [datetime]
     try:
         import cftime
diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
diff --git a/xarray/tests/test_plugins.py b/xarray/tests/test_plugins.py
