implement the PintIndex

pint_xarray/accessors.py

@@ -152,10 +152,9 @@ def __getitem__(self, indexers):
             raise NotImplementedError("pandas-style indexing is not supported, yet")
 
         dims = self.ds.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # convert the indexes to the indexer's units
@@ -165,10 +164,7 @@ def __getitem__(self, indexers):
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         return converted.loc[stripped_indexers]
 
 
@@ -183,10 +179,9 @@ def __getitem__(self, indexers):
             raise NotImplementedError("pandas-style indexing is not supported, yet")
 
         dims = self.da.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # convert the indexes to the indexer's units
@@ -196,10 +191,7 @@ def __getitem__(self, indexers):
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         return converted.loc[stripped_indexers]
 
     def __setitem__(self, indexers, values):
@@ -219,10 +211,7 @@ def __setitem__(self, indexers, values):
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in converted.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(converted)
         self.da.loc[stripped_indexers] = values
 
 
@@ -642,10 +631,9 @@ def reindex(
         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex")
 
         dims = self.da.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # TODO: handle tolerance
@@ -655,10 +643,7 @@ def reindex(
         converted = conversion.convert_units(self.da, indexer_units)
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         indexed = converted.reindex(
             stripped_indexers,
             method=method,
@@ -725,10 +710,9 @@ def interp(
         indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")
 
         dims = self.da.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # convert the indexes to the indexer's units
@@ -737,10 +721,7 @@ def interp(
         stripped = conversion.strip_units(converted)
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         interpolated = stripped.interp(
             stripped_indexers,
             method=method,
@@ -798,10 +779,9 @@ def sel(
         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")
 
         dims = self.da.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # TODO: handle tolerance
@@ -813,10 +793,7 @@ def sel(
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         indexed = converted.sel(
             stripped_indexers,
             method=method,
@@ -866,10 +843,7 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in converted_indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(converted_indexers)
         indexed = self.da.drop_sel(
             stripped_indexers,
             errors=errors,
@@ -1413,10 +1387,9 @@ def reindex(
         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex")
 
         dims = self.ds.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # TODO: handle tolerance
@@ -1426,10 +1399,7 @@ def reindex(
         converted = conversion.convert_units(self.ds, indexer_units)
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         indexed = converted.reindex(
             stripped_indexers,
             method=method,
@@ -1496,10 +1466,9 @@ def interp(
         indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")
 
         dims = self.ds.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # convert the indexes to the indexer's units
@@ -1508,10 +1477,7 @@ def interp(
         stripped = conversion.strip_units(converted)
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         interpolated = stripped.interp(
             stripped_indexers,
             method=method,
@@ -1569,10 +1535,9 @@ def sel(
         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")
 
         dims = self.ds.dims
+        indexer_units = conversion.extract_indexer_units(indexers)
         indexer_units = {
-            name: conversion.extract_indexer_units(indexer)
-            for name, indexer in indexers.items()
-            if name in dims
+            name: indexer for name, indexer in indexer_units.items() if name in dims
         }
 
         # TODO: handle tolerance
@@ -1584,10 +1549,7 @@ def sel(
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(indexers)
         indexed = converted.sel(
             stripped_indexers,
             method=method,
@@ -1639,10 +1601,7 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
             raise KeyError(*e.args) from e
 
         # index
-        stripped_indexers = {
-            name: conversion.strip_indexer_units(indexer)
-            for name, indexer in converted_indexers.items()
-        }
+        stripped_indexers = conversion.strip_indexer_units(converted_indexers)
         indexed = self.ds.drop_sel(
             stripped_indexers,
             errors=errors,

pint_xarray/conversion.py

@@ -1,10 +1,11 @@
 import itertools
 
 import pint
-from xarray import DataArray, Dataset, IndexVariable, Variable
+from xarray import Coordinates, DataArray, Dataset, IndexVariable, Variable
 
 from .compat import call_on_dataset
 from .errors import format_error_message
+from .index import PintMetaIndex
 
 no_unit_values = ("none", None)
 unit_attribute_name = "units"
@@ -122,18 +123,36 @@ def dataset_from_variables(variables, coords, attrs):
 def attach_units_dataset(obj, units):
     attached = {}
     rejected_vars = {}
+
+    indexed_variables = obj.xindexes.variables
     for name, var in obj.variables.items():
+        if name in indexed_variables:
+            continue
+
         unit = units.get(name)
         try:
             converted = attach_units_variable(var, unit)
             attached[name] = converted
         except ValueError as e:
             rejected_vars[name] = (unit, e)
 
+    ds_xindexes = obj.xindexes
+    new_indexes, new_index_vars = ds_xindexes.copy_indexes()
+
+    for idx, idx_vars in ds_xindexes.group_by_index():
+        idx_units = {name: units.get(name) for name in idx_vars.keys()}
+        new_idx = PintMetaIndex(index=idx, units=idx_units)
+        new_indexes.update({k: new_idx for k in idx_vars})
+        new_index_vars.update(new_idx.create_variables(idx_vars))
+
+    new_coords = Coordinates(new_index_vars, new_indexes)
+
     if rejected_vars:
         raise ValueError(rejected_vars)
 
-    return dataset_from_variables(attached, obj._coord_names, obj.attrs)
+    return dataset_from_variables(attached, obj._coord_names, obj.attrs).assign_coords(
+        new_coords
+    )
 
 
 def attach_units(obj, units):
@@ -387,25 +406,31 @@ def convert(indexer, units):
     return converted
 
 
-def extract_indexer_units(indexer):
-    if isinstance(indexer, slice):
-        return slice_extract_units(indexer)
-    elif isinstance(indexer, (DataArray, Variable)):
-        return array_extract_units(indexer.data)
-    else:
-        return array_extract_units(indexer)
+def extract_indexer_units(indexers):
+    def extract(indexer):
+        if isinstance(indexer, slice):
+            return slice_extract_units(indexer)
+        elif isinstance(indexer, (DataArray, Variable)):
+            return array_extract_units(indexer.data)
+        else:
+            return array_extract_units(indexer)
 
+    return {name: extract(indexer) for name, indexer in indexers.items()}
 
-def strip_indexer_units(indexer):
-    if isinstance(indexer, slice):
-        return slice(
-            array_strip_units(indexer.start),
-            array_strip_units(indexer.stop),
-            array_strip_units(indexer.step),
-        )
-    elif isinstance(indexer, DataArray):
-        return strip_units(indexer)
-    elif isinstance(indexer, Variable):
-        return strip_units_variable(indexer)
-    else:
-        return array_strip_units(indexer)
+
+def strip_indexer_units(indexers):
+    def strip(indexer):
+        if isinstance(indexer, slice):
+            return slice(
+                array_strip_units(indexer.start),
+                array_strip_units(indexer.stop),
+                array_strip_units(indexer.step),
+            )
+        elif isinstance(indexer, DataArray):
+            return strip_units(indexer)
+        elif isinstance(indexer, Variable):
+            return strip_units_variable(indexer)
+        else:
+            return array_strip_units(indexer)
+
+    return {name: strip(indexer) for name, indexer in indexers.items()}

pint_xarray/index.py

@@ -0,0 +1,43 @@
+from xarray import Variable
+from xarray.core.indexes import Index, PandasIndex
+
+from . import conversion
+
+
+class PintMetaIndex(Index):
+    # TODO: inherit from MetaIndex once that exists
+    def __init__(self, *, index, units):
+        """create a unit-aware MetaIndex
+
+        Parameters
+        ----------
+        index : xarray.Index
+            The wrapped index object.
+        units : mapping of hashable to unit-like
+            The units of the indexed coordinates
+        """
+        self.index = index
+        self.units = units
+
+    def create_variables(self, variables=None):
+        index_vars = self.index.create_variables(variables)
+
+        index_vars_units = {}
+        for name, var in index_vars.items():
+            data = conversion.array_attach_units(var.data, self.units[name])
+            var_units = Variable(var.dims, data, attrs=var.attrs, encoding=var.encoding)
+            index_vars_units[name] = var_units
+
+        return index_vars_units
+
+    @classmethod
+    def from_variables(cls, variables, options):
+        index = PandasIndex.from_variables(variables)
+        units_dict = {index.index.name: options.get("units")}
+        return cls(index, units_dict)
+
+    def sel(self, labels):
+        converted_labels = conversion.convert_indexer_units(labels, self.units)
+        stripped_labels = conversion.strip_indexer_units(converted_labels)
+
+        return self.index.sel(stripped_labels)