Skip to content

Commit

Permalink
Implement .cf.add_canonical_attributes (#232)
Browse files Browse the repository at this point in the history
  • Loading branch information
malmans2 authored Jun 11, 2021
1 parent a898641 commit f1c35fc
Show file tree
Hide file tree
Showing 10 changed files with 33,719 additions and 33 deletions.
1 change: 1 addition & 0 deletions MANIFEST.in
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
graft cf_xarray/data
9 changes: 2 additions & 7 deletions cf_xarray/__init__.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,5 @@
from pkg_resources import DistributionNotFound, get_distribution

from .accessor import CFAccessor # noqa
from .helpers import bounds_to_vertices, vertices_to_bounds # noqa
from .utils import _get_version

try:
__version__ = get_distribution("cf_xarray").version
except DistributionNotFound:
# package is not installed
__version__ = "unknown"
__version__ = _get_version()
104 changes: 103 additions & 1 deletion cf_xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import re
import warnings
from collections import ChainMap
from datetime import datetime
from typing import (
Any,
Callable,
Expand All @@ -27,10 +28,12 @@
from .criteria import coordinate_criteria, regex
from .helpers import bounds_to_vertices
from .utils import (
_get_version,
_is_datetime_like,
always_iterable,
invert_mappings,
parse_cell_methods_attr,
parse_cf_standard_name_table,
)

#: Classes wrapped by cf_xarray.
Expand Down Expand Up @@ -62,7 +65,6 @@
ATTRS["time"] = ATTRS["T"]
ATTRS["vertical"] = ATTRS["Z"]


# Type for Mapper functions
Mapper = Callable[[Union[DataArray, Dataset], str], List[str]]

Expand Down Expand Up @@ -1521,6 +1523,106 @@ def differentiate(
result *= -1
return result

def add_canonical_attributes(
self,
override: bool = False,
skip: Union[str, List[str]] = None,
verbose: bool = False,
source=None,
) -> Union[Dataset, DataArray]:
"""
Add canonical CF attributes to variables with standard names.
Attributes are parsed from the official CF standard name table.
This function adds an entry to the "history" attribute.
Parameters
----------
override: bool
Override existing attributes.
skip: str, iterable, optional
Attribute(s) to skip: ``{"units", "grib", "amip", "description"}``.
verbose: bool
Print added attributes to screen.
source: optional
Path of `cf-standard-name-table.xml` or file object containing XML data.
If ``None``, use the default version associated with ``cf-xarray``.
Returns
-------
DataArray or Dataset with attributes added.
Notes
-----
The ``"units"`` attribute is never added to datetime-like variables.
"""

# Arguments to add to history
args = ", ".join([f"{k!s}={v!r}" for k, v in locals().items() if k != "self"])

# Defaults
skip = [skip] if isinstance(skip, str) else (skip or [])

# Parse table
info, table, aliases = parse_cf_standard_name_table(source)

# Loop over standard names
ds = self._maybe_to_dataset().copy()
attrs_to_print: dict = {}
for std_name, var_names in ds.cf.standard_names.items():

# Loop over variable names
for var_name in var_names:
old_attrs = ds[var_name].attrs
std_name = aliases.get(std_name, std_name)
new_attrs = table.get(std_name, {})

# Loop over attributes
for key, value in new_attrs.items():
if value and key not in skip and (override or key not in old_attrs):

# Don't add units to time variables (e.g., datetime64, ...)
if key == "units" and _is_datetime_like(ds[var_name]):
continue

# Add attribute
ds[var_name].attrs[key] = value

# Build verbose dictionary
if verbose:
attrs_to_print.setdefault(var_name, {})
attrs_to_print[var_name][key] = value

if verbose:
# Info
strings = ["CF Standard Name Table info:"]
for key, value in info.items():
strings.append(f"- {key}: {value}")

# Attributes added
strings.append("\nAttributes added:")
for varname, attrs in attrs_to_print.items():
strings.append(f"- {varname}:")
for key, value in attrs.items():
strings.append(f" * {key}: {value}")
strings.append("")

print("\n".join(strings))

# Prepend history
now = datetime.now().ctime()
method_name = inspect.stack()[0][3]
version = _get_version()
table_version = info["version_number"]
history = (
f"{now}:"
f" cf.{method_name}({args})"
f" [cf-xarray {version}, cf-standard-name-table {table_version}]\n"
)
obj = self._maybe_to_dataarray(ds)
obj.attrs["history"] = history + obj.attrs.get("history", "")

return obj


@xr.register_dataset_accessor("cf")
class CFDatasetAccessor(CFAccessor):
Expand Down
Loading

0 comments on commit f1c35fc

Please sign in to comment.