Skip to content

create the context objects passed to custom combine_attrs functions #5668

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions xarray/core/merge.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ class Context:
def __init__(self, func):
self.func = func
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@keewis do you have an idea of what to do for groupby.mean and similar calls?

Copy link
Collaborator Author

@keewis keewis Aug 4, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it should be possible to pass the actual (bound?) function (e.g. ds.groupby(...).mean) instead, which would also allow accessing additional data in the combine_attrs function. Not sure how to get a nice repr for that but I guess that's a minor issue.

Edit: to avoid introspecting stack frames, I guess we could pass getattr(self, "name") (where that makes sense)


def __repr__(self):
return f"Context('{self.func}')"


def broadcast_dimension_size(variables: list[Variable]) -> dict[Hashable, int]:
"""Extract dimension sizes from a dictionary of variables.
Expand Down
21 changes: 16 additions & 5 deletions xarray/core/options.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

import warnings
from typing import TYPE_CHECKING, Literal, TypedDict
from typing import TYPE_CHECKING, Callable, Literal, TypedDict

from .utils import FrozenDict

Expand Down Expand Up @@ -90,7 +90,7 @@ def _positive_integer(value: int) -> bool:
"display_expand_data": lambda choice: choice in [True, False, "default"],
"enable_cftimeindex": lambda value: isinstance(value, bool),
"file_cache_maxsize": _positive_integer,
"keep_attrs": lambda choice: choice in [True, False, "default"],
"keep_attrs": lambda choice: choice in [True, False, "default"] or callable(choice),
"use_bottleneck": lambda value: isinstance(value, bool),
"use_flox": lambda value: isinstance(value, bool),
"warn_for_unclosed_files": lambda value: isinstance(value, bool),
Expand Down Expand Up @@ -130,8 +130,18 @@ def _get_boolean_with_default(option: Options, default: bool) -> bool:
)


def _get_keep_attrs(default: bool) -> bool:
return _get_boolean_with_default("keep_attrs", default)
def _get_keep_attrs(default: bool | Callable) -> bool | Callable:
global_choice = OPTIONS["keep_attrs"]

if global_choice == "default":
return default
elif isinstance(global_choice, bool) or callable(global_choice):
return global_choice
else:
raise ValueError(
"The the global option 'keep_attrs' must be one of"
" True, False or 'default', or a callable."
)


class set_options:
Expand Down Expand Up @@ -192,14 +202,15 @@ class set_options:
global least-recently-usage cached. This should be smaller than
your system's per-process file descriptor limit, e.g.,
``ulimit -n`` on Linux.
keep_attrs : {"default", True, False}
keep_attrs : {"default", True, False} or callable
Whether to keep attributes on xarray Datasets/dataarrays after
operations. Can be

* ``True`` : to always keep attrs
* ``False`` : to always discard attrs
* ``default`` : to use original logic that attrs should only
be kept in unambiguous circumstances
* callable: a function that decides what to do.
use_bottleneck : bool, default: True
Whether to use ``bottleneck`` to accelerate 1D reductions and
1D rolling reduction operations.
Expand Down
10 changes: 9 additions & 1 deletion xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -1818,6 +1818,8 @@ def reduce(
Array with summarized data and the indicated dimension(s)
removed.
"""
from .merge import Context, merge_attrs

if dim == ...:
dim = None
if dim is not None and axis is not None:
Expand Down Expand Up @@ -1860,7 +1862,13 @@ def reduce(

if keep_attrs is None:
keep_attrs = _get_keep_attrs(default=False)
attrs = self._attrs if keep_attrs else None

if isinstance(keep_attrs, bool):
keep_attrs = "override" if keep_attrs else "drop"
Comment on lines +1866 to +1867
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can this go in _get_keep_attrs?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no, because we need to support both of these:

Suggested change
if isinstance(keep_attrs, bool):
keep_attrs = "override" if keep_attrs else "drop"
with xr.set_options(keep_attrs=True):
ds.mean()
ds.mean(keep_attrs=True)

but I can move it to merge_attrs

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

merge_attrs would be a good solution.

We could update _get_keep_attrs to _get_keep_attrs(keep_attrs, default=False) and put all the logic in one place but that's a much bigger change.


attrs = merge_attrs(
[self.attrs], combine_attrs=keep_attrs, context=Context(func.__name__)
)

return Variable(dims, data, attrs=attrs)

Expand Down