diff --git a/xarray/core/merge.py b/xarray/core/merge.py index 6262e031a2c..490c18aed12 100644 --- a/xarray/core/merge.py +++ b/xarray/core/merge.py @@ -70,6 +70,9 @@ class Context: def __init__(self, func): self.func = func + 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. diff --git a/xarray/core/options.py b/xarray/core/options.py index 6798c447a2d..dad5545992c 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -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 @@ -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), @@ -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: @@ -192,7 +202,7 @@ 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 @@ -200,6 +210,7 @@ class set_options: * ``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. diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 2d115ff0ed9..3e19cb3322a 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -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: @@ -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" + + attrs = merge_attrs( + [self.attrs], combine_attrs=keep_attrs, context=Context(func.__name__) + ) return Variable(dims, data, attrs=attrs)