Skip to content

Commit

Permalink
Rewrite dicts like .chunks and .sizes (#42)
Browse files Browse the repository at this point in the history
* Rewrite dicts like .chunks and .sizes

* Fix tests and add fallback behaviour.

key-value pairs that are not "special" are copied over.

* Add docs.

* typo
  • Loading branch information
dcherian authored Jul 6, 2020
1 parent 9d53560 commit bd231d8
Show file tree
Hide file tree
Showing 4 changed files with 120 additions and 4 deletions.
33 changes: 32 additions & 1 deletion cf_xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -348,7 +348,38 @@ def _getattr(
An extra decorator, if necessary. This is used by _CFPlotMethods to set default
kwargs based on CF attributes.
"""
func: Callable = getattr(obj, attr)
attribute: Union[Mapping, Callable] = getattr(obj, attr)

if isinstance(attribute, Mapping):
if not attribute:
return dict(attribute)
# attributes like chunks / sizes
newmap = dict()
unused_keys = set(attribute.keys())
for key in _AXIS_NAMES + _COORD_NAMES:
value = _get_axis_coord(obj, key, error=False, default=None)
unused_keys -= set(value)
if value != [None]:
good_values = set(value) & set(obj.dims)
if not good_values:
continue
if len(good_values) > 1:
raise AttributeError(
f"cf_xarray can't wrap attribute {attr!r} because there are multiple values for {key!r} viz. {good_values!r}. "
f"There is no unique mapping from {key!r} to a value in {attr!r}."
)
newmap.update({key: attribute[good_values.pop()]})
newmap.update({key: attribute[key] for key in unused_keys})
return newmap

elif isinstance(attribute, Callable): # type: ignore
func: Callable = attribute

else:
raise AttributeError(
f"cf_xarray does not know how to wrap attribute '{type(obj).__name__}.{attr}'. "
"Please file an issue if you have a solution."
)

@functools.wraps(func)
def wrapper(*args, **kwargs):
Expand Down
1 change: 1 addition & 0 deletions cf_xarray/tests/datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
)
anc


multiple = xr.Dataset()
multiple.coords["x1"] = ("x1", range(30), {"axis": "X"})
multiple.coords["y1"] = ("y1", range(20), {"axis": "Y"})
Expand Down
34 changes: 34 additions & 0 deletions cf_xarray/tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,40 @@
objects = datasets + dataarrays


def test_dicts():
from .datasets import airds

actual = airds.cf.sizes
expected = {"X": 50, "Y": 25, "T": 4, "longitude": 50, "latitude": 25, "time": 4}
assert actual == expected

assert popds.cf.sizes == popds.sizes

with pytest.raises(AttributeError):
multiple.cf.sizes

assert airds.cf.chunks == {}

expected = {
"X": (50,),
"Y": (5, 5, 5, 5, 5),
"T": (4,),
"longitude": (50,),
"latitude": (5, 5, 5, 5, 5),
"time": (4,),
}
assert airds.chunk({"lat": 5}).cf.chunks == expected

with pytest.raises(AttributeError):
airds.da.cf.chunks

airds = airds.copy(deep=True)
airds.lon.attrs = {}
actual = airds.cf.sizes
expected = {"lon": 50, "Y": 25, "T": 4, "latitude": 25, "time": 4}
assert actual == expected


def test_describe():
actual = airds.cf._describe()
expected = (
Expand Down
56 changes: 53 additions & 3 deletions doc/examples/introduction.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,58 @@
"pop.cf.get_valid_keys()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Feature: Rewriting property dictionaries\n",
"\n",
"`cf_xarray` will rewrite the `.sizes` and `.chunks` dictionaries so that one can index by a special CF axis or coordinate name"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ds.cf.sizes"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Note the duplicate entries above:\n",
"\n",
"1. One for `X`, `Y`, `T`\n",
"2. and one for `longitude`, `latitude` and `time`.\n",
"\n",
"An error is raised if there are multiple `'X'` variables (for example)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"raises-exception"
]
},
"outputs": [],
"source": [
"multiple.cf.sizes"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"multiple.v1.cf.sizes"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down Expand Up @@ -637,9 +689,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"#### miscellaneous features\n",
"\n",
"You can mix \"special names\" and variable names"
"## Feature: mix \"special names\" and variable names"
]
},
{
Expand Down

0 comments on commit bd231d8

Please sign in to comment.