diff --git a/setup.cfg b/setup.cfg index c80ff300a60..18922b1647a 100644 --- a/setup.cfg +++ b/setup.cfg @@ -45,6 +45,8 @@ ignore_missing_imports = True ignore_missing_imports = True [mypy-Nio.*] ignore_missing_imports = True +[mypy-nc_time_axis.*] +ignore_missing_imports = True [mypy-numpy.*] ignore_missing_imports = True [mypy-netCDF4.*] diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 633e695694b..f9fd9022de9 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -7,8 +7,8 @@ from collections import Counter, OrderedDict from distutils.version import LooseVersion from typing import ( - AbstractSet, Any, Dict, Iterable, List, Mapping, Union, Tuple, - TYPE_CHECKING, TypeVar + AbstractSet, Any, Callable, Iterable, List, Mapping, Optional, Sequence, + Tuple, TYPE_CHECKING, Union, ) import numpy as np @@ -190,20 +190,19 @@ def build_output_coords( return output_coords -def apply_dataarray_ufunc(func, *args, **kwargs): - """apply_dataarray_ufunc(func, *args, signature, join='inner', - exclude_dims=frozenset()) +def apply_dataarray_vfunc( + func, + *args, + signature, + join='inner', + exclude_dims=frozenset(), + keep_attrs=False +): + """Apply a variable level function over DataArray, Variable and/or ndarray + objects. """ from .dataarray import DataArray - signature = kwargs.pop('signature') - join = kwargs.pop('join', 'inner') - exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET) - keep_attrs = kwargs.pop('keep_attrs', True) - if kwargs: - raise TypeError('apply_dataarray_ufunc() got unexpected keyword ' - 'arguments: %s' % list(kwargs)) - if len(args) > 1: args = deep_align(args, join=join, copy=False, exclude=exclude_dims, raise_on_invalid=False) @@ -261,15 +260,19 @@ def assert_and_return_exact_match(all_keys): } -def join_dict_keys(objects, how='inner'): - # type: (Iterable[Union[Mapping, Any]], str) -> Iterable +def join_dict_keys( + objects: Iterable[Union[Mapping, Any]], how: str = 'inner', +) -> Iterable: joiner = _JOINERS[how] all_keys = [obj.keys() for obj in objects if hasattr(obj, 'keys')] return joiner(all_keys) -def collect_dict_values(objects, keys, fill_value=None): - # type: (Iterable[Union[Mapping, Any]], Iterable, Any) -> List[list] +def collect_dict_values( + objects: Iterable[Union[Mapping, Any]], + keys: Iterable, + fill_value: object = None, +) -> List[list]: return [[obj.get(key, fill_value) if is_dict_like(obj) else obj @@ -298,17 +301,12 @@ def _unpack_dict_tuples( return out -def apply_dict_of_variables_ufunc(func, *args, **kwargs): - """apply_dict_of_variables_ufunc(func, *args, signature, join='inner', - fill_value=None): +def apply_dict_of_variables_vfunc( + func, *args, signature, join='inner', fill_value=None +): + """Apply a variable level function over dicts of DataArray, DataArray, + Variable and ndarray objects. """ - signature = kwargs.pop('signature') - join = kwargs.pop('join', 'inner') - fill_value = kwargs.pop('fill_value', None) - if kwargs: - raise TypeError('apply_dict_of_variables_ufunc() got unexpected ' - 'keyword arguments: %s' % list(kwargs)) - args = [_as_variables_or_variable(arg) for arg in args] names = join_dict_keys(args, how=join) grouped_by_name = collect_dict_values(args, names, fill_value) @@ -323,8 +321,10 @@ def apply_dict_of_variables_ufunc(func, *args, **kwargs): return result_vars -def _fast_dataset(variables, coord_variables): - # type: (OrderedDict[Any, Variable], Mapping[Any, Variable]) -> Dataset +def _fast_dataset( + variables: 'OrderedDict[Any, Variable]', + coord_variables: Mapping[Any, Variable], +) -> 'Dataset': """Create a dataset as quickly as possible. Beware: the `variables` OrderedDict is modified INPLACE. @@ -335,21 +335,20 @@ def _fast_dataset(variables, coord_variables): return Dataset._from_vars_and_coord_names(variables, coord_names) -def apply_dataset_ufunc(func, *args, **kwargs): - """apply_dataset_ufunc(func, *args, signature, join='inner', - dataset_join='inner', fill_value=None, - exclude_dims=frozenset(), keep_attrs=False): - - If dataset_join != 'inner', a non-default fill_value must be supplied - by the user. Otherwise a TypeError is raised. +def apply_dataset_vfunc( + func, + *args, + signature, + join='inner', + dataset_join='exact', + fill_value=_NO_FILL_VALUE, + exclude_dims=frozenset(), + keep_attrs=False +): + """Apply a variable level function over Dataset, dict of DataArray, + DataArray, Variable and/or ndarray objects. """ from .dataset import Dataset - signature = kwargs.pop('signature') - join = kwargs.pop('join', 'inner') - dataset_join = kwargs.pop('dataset_join', 'inner') - fill_value = kwargs.pop('fill_value', None) - exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET) - keep_attrs = kwargs.pop('keep_attrs', False) first_obj = args[0] # we'll copy attrs from this in case keep_attrs=True if (dataset_join not in _JOINS_WITHOUT_FILL_VALUES and @@ -358,9 +357,6 @@ def apply_dataset_ufunc(func, *args, **kwargs): 'data variables with apply_ufunc, you must supply the ' 'dataset_fill_value argument.') - if kwargs: - raise TypeError('apply_dataset_ufunc() got unexpected keyword ' - 'arguments: %s' % list(kwargs)) if len(args) > 1: args = deep_align(args, join=join, copy=False, exclude=exclude_dims, raise_on_invalid=False) @@ -368,7 +364,7 @@ def apply_dataset_ufunc(func, *args, **kwargs): list_of_coords = build_output_coords(args, signature, exclude_dims) args = [getattr(arg, 'data_vars', arg) for arg in args] - result_vars = apply_dict_of_variables_ufunc( + result_vars = apply_dict_of_variables_vfunc( func, *args, signature=signature, join=dataset_join, fill_value=fill_value) @@ -402,7 +398,10 @@ def _iter_over_selections(obj, dim, values): yield obj_sel -def apply_groupby_ufunc(func, *args): +def apply_groupby_func(func, *args): + """Apply a dataset or datarray level function over GroupBy, Dataset, + DataArray, Variable and/or ndarray objects. + """ from .groupby import GroupBy, peek_at from .variable import Variable @@ -444,7 +443,7 @@ def apply_groupby_ufunc(func, *args): def unified_dim_sizes( variables: Iterable[Variable], - exclude_dims: AbstractSet = frozenset(), + exclude_dims: AbstractSet = frozenset() ) -> 'OrderedDict[Any, int]': dim_sizes = OrderedDict() # type: OrderedDict[Any, int] @@ -516,21 +515,20 @@ def broadcast_compat_data(variable, broadcast_dims, core_dims): return data -def apply_variable_ufunc(func, *args, **kwargs): - """apply_variable_ufunc(func, *args, signature, exclude_dims=frozenset()) +def apply_variable_ufunc( + func, + *args, + signature, + exclude_dims=frozenset(), + dask='forbidden', + output_dtypes=None, + output_sizes=None, + keep_attrs=False +): + """Apply a ndarray level function over Variable and/or ndarray objects. """ from .variable import Variable, as_compatible_data - signature = kwargs.pop('signature') - exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET) - dask = kwargs.pop('dask', 'forbidden') - output_dtypes = kwargs.pop('output_dtypes', None) - output_sizes = kwargs.pop('output_sizes', None) - keep_attrs = kwargs.pop('keep_attrs', False) - if kwargs: - raise TypeError('apply_variable_ufunc() got unexpected keyword ' - 'arguments: %s' % list(kwargs)) - dim_sizes = unified_dim_sizes((a for a in args if hasattr(a, 'dims')), exclude_dims=exclude_dims) broadcast_dims = tuple(dim for dim in dim_sizes @@ -661,14 +659,8 @@ def _apply_with_dask_atop(func, args, input_dims, output_dims, signature, new_axes=output_sizes) -def apply_array_ufunc(func, *args, **kwargs): - """apply_array_ufunc(func, *args, dask='forbidden') - """ - dask = kwargs.pop('dask', 'forbidden') - if kwargs: - raise TypeError('apply_array_ufunc() got unexpected keyword ' - 'arguments: %s' % list(kwargs)) - +def apply_array_ufunc(func, *args, dask='forbidden'): + """Apply a ndarray level function over ndarray objects.""" if any(isinstance(arg, dask_array_type) for arg in args): if dask == 'forbidden': raise ValueError('apply_ufunc encountered a dask array on an ' @@ -687,23 +679,23 @@ def apply_array_ufunc(func, *args, **kwargs): return func(*args) -def apply_ufunc(func, *args, **kwargs): - """apply_ufunc(func : Callable, - *args : Any, - input_core_dims : Optional[Sequence[Sequence]] = None, - output_core_dims : Optional[Sequence[Sequence]] = ((),), - exclude_dims : Collection = frozenset(), - vectorize : bool = False, - join : str = 'exact', - dataset_join : str = 'exact', - dataset_fill_value : Any = _NO_FILL_VALUE, - keep_attrs : bool = False, - kwargs : Mapping = None, - dask : str = 'forbidden', - output_dtypes : Optional[Sequence] = None, - output_sizes : Optional[Mapping[Any, int]] = None) - - Apply a vectorized function for unlabeled arrays on xarray objects. +def apply_ufunc( + func: Callable, + *args: Any, + input_core_dims: Optional[Sequence[Sequence]] = None, + output_core_dims: Optional[Sequence[Sequence]] = ((),), + exclude_dims: AbstractSet = frozenset(), + vectorize: bool = False, + join: str = 'exact', + dataset_join: str = 'exact', + dataset_fill_value: object = _NO_FILL_VALUE, + keep_attrs: bool = False, + kwargs: Mapping = None, + dask: str = 'forbidden', + output_dtypes: Optional[Sequence] = None, + output_sizes: Optional[Mapping[Any, int]] = None +) -> Any: + """Apply a vectorized function for unlabeled arrays on xarray objects. The function will be mapped over the data variable(s) of the input arguments using xarray's standard rules for labeled computation, including @@ -907,22 +899,6 @@ def earth_mover_distance(first_samples, from .dataarray import DataArray from .variable import Variable - input_core_dims = kwargs.pop('input_core_dims', None) - output_core_dims = kwargs.pop('output_core_dims', ((),)) - vectorize = kwargs.pop('vectorize', False) - join = kwargs.pop('join', 'exact') - dataset_join = kwargs.pop('dataset_join', 'exact') - keep_attrs = kwargs.pop('keep_attrs', False) - exclude_dims = kwargs.pop('exclude_dims', frozenset()) - dataset_fill_value = kwargs.pop('dataset_fill_value', _NO_FILL_VALUE) - kwargs_ = kwargs.pop('kwargs', None) - dask = kwargs.pop('dask', 'forbidden') - output_dtypes = kwargs.pop('output_dtypes', None) - output_sizes = kwargs.pop('output_sizes', None) - if kwargs: - raise TypeError('apply_ufunc() got unexpected keyword arguments: %s' - % list(kwargs)) - if input_core_dims is None: input_core_dims = ((),) * (len(args)) elif len(input_core_dims) != len(args): @@ -931,14 +907,17 @@ def earth_mover_distance(first_samples, 'the number of arguments. Given input_core_dims: {}, ' 'number of args: {}.'.format(input_core_dims, len(args))) + if kwargs is None: + kwargs = {} + signature = _UFuncSignature(input_core_dims, output_core_dims) if exclude_dims and not exclude_dims <= signature.all_core_dims: raise ValueError('each dimension in `exclude_dims` must also be a ' 'core dimension in the function signature') - if kwargs_: - func = functools.partial(func, **kwargs_) + if kwargs: + func = functools.partial(func, **kwargs) if vectorize: if signature.all_core_dims: @@ -950,14 +929,11 @@ def earth_mover_distance(first_samples, 'dimensions.') func = np.vectorize(func, otypes=output_dtypes, - signature=signature.to_gufunc_string(), - excluded=set(kwargs)) + signature=signature.to_gufunc_string()) else: - func = np.vectorize(func, - otypes=output_dtypes, - excluded=set(kwargs)) + func = np.vectorize(func, otypes=output_dtypes) - variables_ufunc = functools.partial(apply_variable_ufunc, func, + variables_vfunc = functools.partial(apply_variable_ufunc, func, signature=signature, exclude_dims=exclude_dims, keep_attrs=keep_attrs, @@ -966,7 +942,6 @@ def earth_mover_distance(first_samples, output_sizes=output_sizes) if any(isinstance(a, GroupBy) for a in args): - # kwargs has already been added into func this_apply = functools.partial(apply_ufunc, func, input_core_dims=input_core_dims, output_core_dims=output_core_dims, @@ -976,31 +951,29 @@ def earth_mover_distance(first_samples, dataset_fill_value=dataset_fill_value, keep_attrs=keep_attrs, dask=dask) - return apply_groupby_ufunc(this_apply, *args) + return apply_groupby_func(this_apply, *args) elif any(is_dict_like(a) for a in args): - return apply_dataset_ufunc(variables_ufunc, *args, + return apply_dataset_vfunc(variables_vfunc, *args, signature=signature, join=join, exclude_dims=exclude_dims, - fill_value=dataset_fill_value, dataset_join=dataset_join, + fill_value=dataset_fill_value, keep_attrs=keep_attrs) elif any(isinstance(a, DataArray) for a in args): - return apply_dataarray_ufunc(variables_ufunc, *args, + return apply_dataarray_vfunc(variables_vfunc, *args, signature=signature, join=join, exclude_dims=exclude_dims, keep_attrs=keep_attrs) elif any(isinstance(a, Variable) for a in args): - return variables_ufunc(*args) + return variables_vfunc(*args) else: return apply_array_ufunc(func, *args, dask=dask) -def dot(*arrays, **kwargs): - """ dot(*arrays, dims=None) - - Generalized dot product for xarray objects. Like np.einsum, but +def dot(*arrays, dims=None, **kwargs): + """Generalized dot product for xarray objects. Like np.einsum, but provides a simpler interface based on array dimensions. Parameters @@ -1036,8 +1009,6 @@ def dot(*arrays, **kwargs): from .dataarray import DataArray from .variable import Variable - dims = kwargs.pop('dims', None) - if any(not isinstance(arr, (Variable, DataArray)) for arr in arrays): raise TypeError('Only xr.DataArray and xr.Variable are supported.' 'Given {}.'.format([type(arr) for arr in arrays]))