From 22c2c023d52f8f3e52b53c1e1c77554e07c37170 Mon Sep 17 00:00:00 2001 From: Alan Brammer Date: Wed, 5 Apr 2023 22:01:35 +0000 Subject: [PATCH 01/17] Initial commit --- xarray/core/_typed_ops.py | 64 +++++++++++++++++++++++++++++++++++++ xarray/core/ops.py | 2 ++ xarray/util/generate_ops.py | 4 +++ 3 files changed, 70 insertions(+) diff --git a/xarray/core/_typed_ops.py b/xarray/core/_typed_ops.py index a6e6fdbfaec..d3a783be45d 100644 --- a/xarray/core/_typed_ops.py +++ b/xarray/core/_typed_ops.py @@ -42,6 +42,12 @@ def __xor__(self, other): def __or__(self, other): return self._binary_op(other, operator.or_) + def __lshift__(self, other): + return self._binary_op(other, operator.lshift) + + def __rshift__(self, other): + return self._binary_op(other, operator.rshift) + def __lt__(self, other): return self._binary_op(other, operator.lt) @@ -123,6 +129,12 @@ def __ixor__(self, other): def __ior__(self, other): return self._inplace_binary_op(other, operator.ior) + def __ilshift__(self, other): + return self._inplace_binary_op(other, operator.ilshift) + + def __irshift__(self, other): + return self._inplace_binary_op(other, operator.irshift) + def _unary_op(self, f, *args, **kwargs): raise NotImplementedError @@ -160,6 +172,8 @@ def conjugate(self, *args, **kwargs): __and__.__doc__ = operator.and_.__doc__ __xor__.__doc__ = operator.xor.__doc__ __or__.__doc__ = operator.or_.__doc__ + __lshift__.__doc__ = operator.lshift.__doc__ + __rshift__.__doc__ = operator.rshift.__doc__ __lt__.__doc__ = operator.lt.__doc__ __le__.__doc__ = operator.le.__doc__ __gt__.__doc__ = operator.gt.__doc__ @@ -186,6 +200,8 @@ def conjugate(self, *args, **kwargs): __iand__.__doc__ = operator.iand.__doc__ __ixor__.__doc__ = operator.ixor.__doc__ __ior__.__doc__ = operator.ior.__doc__ + __ilshift__.__doc__ = operator.ilshift.__doc__ + __irshift__.__doc__ = operator.irshift.__doc__ __neg__.__doc__ = operator.neg.__doc__ __pos__.__doc__ = operator.pos.__doc__ __abs__.__doc__ = operator.abs.__doc__ @@ -232,6 +248,12 @@ def __xor__(self, other): def __or__(self, other): return self._binary_op(other, operator.or_) + def __lshift__(self, other): + return self._binary_op(other, operator.lshift) + + def __rshift__(self, other): + return self._binary_op(other, operator.rshift) + def __lt__(self, other): return self._binary_op(other, operator.lt) @@ -313,6 +335,12 @@ def __ixor__(self, other): def __ior__(self, other): return self._inplace_binary_op(other, operator.ior) + def __ilshift__(self, other): + return self._inplace_binary_op(other, operator.ilshift) + + def __irshift__(self, other): + return self._inplace_binary_op(other, operator.irshift) + def _unary_op(self, f, *args, **kwargs): raise NotImplementedError @@ -350,6 +378,8 @@ def conjugate(self, *args, **kwargs): __and__.__doc__ = operator.and_.__doc__ __xor__.__doc__ = operator.xor.__doc__ __or__.__doc__ = operator.or_.__doc__ + __lshift__.__doc__ = operator.lshift.__doc__ + __rshift__.__doc__ = operator.rshift.__doc__ __lt__.__doc__ = operator.lt.__doc__ __le__.__doc__ = operator.le.__doc__ __gt__.__doc__ = operator.gt.__doc__ @@ -376,6 +406,8 @@ def conjugate(self, *args, **kwargs): __iand__.__doc__ = operator.iand.__doc__ __ixor__.__doc__ = operator.ixor.__doc__ __ior__.__doc__ = operator.ior.__doc__ + __ilshift__.__doc__ = operator.ilshift.__doc__ + __irshift__.__doc__ = operator.irshift.__doc__ __neg__.__doc__ = operator.neg.__doc__ __pos__.__doc__ = operator.pos.__doc__ __abs__.__doc__ = operator.abs.__doc__ @@ -422,6 +454,12 @@ def __xor__(self, other): def __or__(self, other): return self._binary_op(other, operator.or_) + def __lshift__(self, other): + return self._binary_op(other, operator.lshift) + + def __rshift__(self, other): + return self._binary_op(other, operator.rshift) + def __lt__(self, other): return self._binary_op(other, operator.lt) @@ -503,6 +541,12 @@ def __ixor__(self, other): def __ior__(self, other): return self._inplace_binary_op(other, operator.ior) + def __ilshift__(self, other): + return self._inplace_binary_op(other, operator.ilshift) + + def __irshift__(self, other): + return self._inplace_binary_op(other, operator.irshift) + def _unary_op(self, f, *args, **kwargs): raise NotImplementedError @@ -540,6 +584,8 @@ def conjugate(self, *args, **kwargs): __and__.__doc__ = operator.and_.__doc__ __xor__.__doc__ = operator.xor.__doc__ __or__.__doc__ = operator.or_.__doc__ + __lshift__.__doc__ = operator.lshift.__doc__ + __rshift__.__doc__ = operator.rshift.__doc__ __lt__.__doc__ = operator.lt.__doc__ __le__.__doc__ = operator.le.__doc__ __gt__.__doc__ = operator.gt.__doc__ @@ -566,6 +612,8 @@ def conjugate(self, *args, **kwargs): __iand__.__doc__ = operator.iand.__doc__ __ixor__.__doc__ = operator.ixor.__doc__ __ior__.__doc__ = operator.ior.__doc__ + __ilshift__.__doc__ = operator.ilshift.__doc__ + __irshift__.__doc__ = operator.irshift.__doc__ __neg__.__doc__ = operator.neg.__doc__ __pos__.__doc__ = operator.pos.__doc__ __abs__.__doc__ = operator.abs.__doc__ @@ -612,6 +660,12 @@ def __xor__(self, other): def __or__(self, other): return self._binary_op(other, operator.or_) + def __lshift__(self, other): + return self._binary_op(other, operator.lshift) + + def __rshift__(self, other): + return self._binary_op(other, operator.rshift) + def __lt__(self, other): return self._binary_op(other, operator.lt) @@ -670,6 +724,8 @@ def __ror__(self, other): __and__.__doc__ = operator.and_.__doc__ __xor__.__doc__ = operator.xor.__doc__ __or__.__doc__ = operator.or_.__doc__ + __lshift__.__doc__ = operator.lshift.__doc__ + __rshift__.__doc__ = operator.rshift.__doc__ __lt__.__doc__ = operator.lt.__doc__ __le__.__doc__ = operator.le.__doc__ __gt__.__doc__ = operator.gt.__doc__ @@ -724,6 +780,12 @@ def __xor__(self, other): def __or__(self, other): return self._binary_op(other, operator.or_) + def __lshift__(self, other): + return self._binary_op(other, operator.lshift) + + def __rshift__(self, other): + return self._binary_op(other, operator.rshift) + def __lt__(self, other): return self._binary_op(other, operator.lt) @@ -782,6 +844,8 @@ def __ror__(self, other): __and__.__doc__ = operator.and_.__doc__ __xor__.__doc__ = operator.xor.__doc__ __or__.__doc__ = operator.or_.__doc__ + __lshift__.__doc__ = operator.lshift.__doc__ + __rshift__.__doc__ = operator.rshift.__doc__ __lt__.__doc__ = operator.lt.__doc__ __le__.__doc__ = operator.le.__doc__ __gt__.__doc__ = operator.gt.__doc__ diff --git a/xarray/core/ops.py b/xarray/core/ops.py index 009616c5f12..e1c3573841a 100644 --- a/xarray/core/ops.py +++ b/xarray/core/ops.py @@ -33,6 +33,8 @@ "and", "xor", "or", + "lshift", + "rshift", ] # methods which pass on the numpy return value unchanged diff --git a/xarray/util/generate_ops.py b/xarray/util/generate_ops.py index 02a3725f475..cf0673e7cca 100644 --- a/xarray/util/generate_ops.py +++ b/xarray/util/generate_ops.py @@ -30,6 +30,8 @@ ("__and__", "operator.and_"), ("__xor__", "operator.xor"), ("__or__", "operator.or_"), + ("__lshift__", "operator.lshift"), + ("__rshift__", "operator.rshift"), ) BINOPS_REFLEXIVE = ( ("__radd__", "operator.add"), @@ -54,6 +56,8 @@ ("__iand__", "operator.iand"), ("__ixor__", "operator.ixor"), ("__ior__", "operator.ior"), + ("__ilshift__", "operator.ilshift"), + ("__irshift__", "operator.irshift"), ) UNARY_OPS = ( ("__neg__", "operator.neg"), From 29dd9178e1fec627eedd3b5eb43631016d2b3313 Mon Sep 17 00:00:00 2001 From: Alan Brammer Date: Thu, 6 Apr 2023 03:21:46 +0000 Subject: [PATCH 02/17] Add auto generated .pyi typed_ops file --- xarray/core/_typed_ops.pyi | 50 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/xarray/core/_typed_ops.pyi b/xarray/core/_typed_ops.pyi index 98a17a47cd5..9e2ba2d3a06 100644 --- a/xarray/core/_typed_ops.pyi +++ b/xarray/core/_typed_ops.pyi @@ -44,6 +44,8 @@ class DatasetOpsMixin: def __and__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... def __xor__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... def __or__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... + def __lshift__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... + def __rshift__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... def __lt__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... def __le__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... def __gt__(self: T_Dataset, other: DsCompatible) -> T_Dataset: ... @@ -135,6 +137,18 @@ class DataArrayOpsMixin: @overload def __or__(self: T_DataArray, other: DaCompatible) -> T_DataArray: ... @overload + def __lshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __lshift__(self, other: "DatasetGroupBy") -> "Dataset": ... + @overload + def __lshift__(self: T_DataArray, other: DaCompatible) -> T_DataArray: ... + @overload + def __rshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __rshift__(self, other: "DatasetGroupBy") -> "Dataset": ... + @overload + def __rshift__(self: T_DataArray, other: DaCompatible) -> T_DataArray: ... + @overload def __lt__(self, other: T_Dataset) -> T_Dataset: ... @overload def __lt__(self, other: "DatasetGroupBy") -> "Dataset": ... @@ -305,6 +319,18 @@ class VariableOpsMixin: @overload def __or__(self: T_Variable, other: VarCompatible) -> T_Variable: ... @overload + def __lshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __lshift__(self, other: T_DataArray) -> T_DataArray: ... + @overload + def __lshift__(self: T_Variable, other: VarCompatible) -> T_Variable: ... + @overload + def __rshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __rshift__(self, other: T_DataArray) -> T_DataArray: ... + @overload + def __rshift__(self: T_Variable, other: VarCompatible) -> T_Variable: ... + @overload def __lt__(self, other: T_Dataset) -> T_Dataset: ... @overload def __lt__(self, other: T_DataArray) -> T_DataArray: ... @@ -475,6 +501,18 @@ class DatasetGroupByOpsMixin: @overload def __or__(self, other: GroupByIncompatible) -> NoReturn: ... @overload + def __lshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __lshift__(self, other: "DataArray") -> "Dataset": ... + @overload + def __lshift__(self, other: GroupByIncompatible) -> NoReturn: ... + @overload + def __rshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __rshift__(self, other: "DataArray") -> "Dataset": ... + @overload + def __rshift__(self, other: GroupByIncompatible) -> NoReturn: ... + @overload def __lt__(self, other: T_Dataset) -> T_Dataset: ... @overload def __lt__(self, other: "DataArray") -> "Dataset": ... @@ -635,6 +673,18 @@ class DataArrayGroupByOpsMixin: @overload def __or__(self, other: GroupByIncompatible) -> NoReturn: ... @overload + def __lshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __lshift__(self, other: T_DataArray) -> T_DataArray: ... + @overload + def __lshift__(self, other: GroupByIncompatible) -> NoReturn: ... + @overload + def __rshift__(self, other: T_Dataset) -> T_Dataset: ... + @overload + def __rshift__(self, other: T_DataArray) -> T_DataArray: ... + @overload + def __rshift__(self, other: GroupByIncompatible) -> NoReturn: ... + @overload def __lt__(self, other: T_Dataset) -> T_Dataset: ... @overload def __lt__(self, other: T_DataArray) -> T_DataArray: ... From e14da23a5ca3f9ea371c67d030cd606ddbfc8918 Mon Sep 17 00:00:00 2001 From: abrammer Date: Thu, 20 Apr 2023 20:48:23 -0600 Subject: [PATCH 03/17] Add bitshift op test for dask --- xarray/tests/test_dask.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index b37399d6ef8..ad48000328b 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -178,6 +178,18 @@ def test_binary_op(self): self.assertLazyAndIdentical(u + u, v + v) self.assertLazyAndIdentical(u[0] + u, v[0] + v) + def test_binary_op_bitshift(self): + # bit shifts only work on ints so we need to generate + # new eager and lazy vars + values = np.random.default_rng(0).integers(0, high=1e10, size=(4, 6), dtype=np.int64) + data = da.from_array(values, chunks=(2, 2)) + u = Variable(("x", "y"), values) + v = Variable(("x", "y"), data) + self.assertLazyAndIdentical(u << 2, v << 2) + self.assertLazyAndIdentical(u << 5, v << 5) + self.assertLazyAndIdentical(u >> 2, v >> 2) + self.assertLazyAndIdentical(u >> 5, v >> 5) + def test_repr(self): expected = dedent( """\ From 426d559e3528003aff788ec0405141e4e577ebe6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Apr 2023 02:49:38 +0000 Subject: [PATCH 04/17] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_dask.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index ad48000328b..2774c97ee10 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -179,9 +179,11 @@ def test_binary_op(self): self.assertLazyAndIdentical(u[0] + u, v[0] + v) def test_binary_op_bitshift(self): - # bit shifts only work on ints so we need to generate + # bit shifts only work on ints so we need to generate # new eager and lazy vars - values = np.random.default_rng(0).integers(0, high=1e10, size=(4, 6), dtype=np.int64) + values = np.random.default_rng(0).integers( + 0, high=1e10, size=(4, 6), dtype=np.int64 + ) data = da.from_array(values, chunks=(2, 2)) u = Variable(("x", "y"), values) v = Variable(("x", "y"), data) From 5a78c53cfd732cba0bcb3439ec662d83b78d60b7 Mon Sep 17 00:00:00 2001 From: abrammer Date: Thu, 20 Apr 2023 21:31:36 -0600 Subject: [PATCH 05/17] Add bitshift tests to dataarray and variable --- xarray/tests/test_dataarray.py | 5 +++++ xarray/tests/test_variable.py | 22 +++++++++++++++------- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 4ffa95e16e6..a8caf3f6d82 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -3905,6 +3905,11 @@ def test_binary_op_propagate_indexes(self) -> None: actual = (self.dv > 10).xindexes["x"] assert expected is actual + # use mda for bitshift test as it's type int + actual = (self.mda << 2).xindexes["x"] + expected = self.mda.xindexes["x"] + assert expected is actual + def test_binary_op_join_setting(self) -> None: dim = "x" align_type: Final = "outer" diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 2f571eb9a83..c6df7bcea44 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -335,9 +335,10 @@ def test_pandas_period_index(self): assert v[0] == pd.Period("2000", freq="B") assert "Period('2000-01-03', 'B')" in repr(v) - def test_1d_math(self): - x = 1.0 * np.arange(5) - y = np.ones(5) + @pytest.mark.parametrize("dtype", [float, int]) + def test_1d_math(self, dtype): + x = np.arange(5, dtype=dtype) + y = np.ones(5, dtype=dtype) # should we need `.to_base_variable()`? # probably a break that `+v` changes type? @@ -351,11 +352,18 @@ def test_1d_math(self): assert_identical(base_v, v + 0) assert_identical(base_v, 0 + v) assert_identical(base_v, v * 1) + if dtype is int: + assert_identical(base_v, v << 0) + assert_array_equal(v << 3, x << 3) + assert_array_equal(v >> 2, x >> 2) # binary ops with numpy arrays assert_array_equal((v * x).values, x**2) assert_array_equal((x * v).values, x**2) assert_array_equal(v - y, v - 1) assert_array_equal(y - v, 1 - v) + if dtype is int: + assert_array_equal(v << x, x << x) + assert_array_equal(v >> x, x >> x) # verify attributes are dropped v2 = self.cls(["x"], x, {"units": "meters"}) with set_options(keep_attrs=False): @@ -369,10 +377,10 @@ def test_1d_math(self): # something complicated assert_array_equal((v**2 * w - 1 + x).values, x**2 * y - 1 + x) # make sure dtype is preserved (for Index objects) - assert float == (+v).dtype - assert float == (+v).values.dtype - assert float == (0 + v).dtype - assert float == (0 + v).values.dtype + assert dtype == (+v).dtype + assert dtype == (+v).values.dtype + assert dtype == (0 + v).dtype + assert dtype == (0 + v).values.dtype # check types of returned data assert isinstance(+v, Variable) assert not isinstance(+v, IndexVariable) From 9961aa5851bb32a77c678829db25cbf32eb39e19 Mon Sep 17 00:00:00 2001 From: Alan Brammer Date: Fri, 21 Apr 2023 12:19:10 -0600 Subject: [PATCH 06/17] Apply typing suggestions from code review Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com> --- xarray/tests/test_dask.py | 2 +- xarray/tests/test_variable.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 2774c97ee10..824d9d0a39d 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -178,7 +178,7 @@ def test_binary_op(self): self.assertLazyAndIdentical(u + u, v + v) self.assertLazyAndIdentical(u[0] + u, v[0] + v) - def test_binary_op_bitshift(self): + def test_binary_op_bitshift(self) -> None: # bit shifts only work on ints so we need to generate # new eager and lazy vars values = np.random.default_rng(0).integers( diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 7fa9c344ba0..238cba86273 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -345,7 +345,7 @@ def test_pandas_period_index(self): assert "Period('2000-01-03', 'B')" in repr(v) @pytest.mark.parametrize("dtype", [float, int]) - def test_1d_math(self, dtype): + def test_1d_math(self, dtype: np.typing.DTypeLike) -> None: x = np.arange(5, dtype=dtype) y = np.ones(5, dtype=dtype) From 13193bc58323a1c35fc4ce0f576b4c94cda103ce Mon Sep 17 00:00:00 2001 From: abrammer Date: Sat, 22 Apr 2023 13:16:26 -0600 Subject: [PATCH 07/17] Fix type checking on test_dask addition --- xarray/tests/test_dask.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 824d9d0a39d..1171464a962 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -181,9 +181,8 @@ def test_binary_op(self): def test_binary_op_bitshift(self) -> None: # bit shifts only work on ints so we need to generate # new eager and lazy vars - values = np.random.default_rng(0).integers( - 0, high=1e10, size=(4, 6), dtype=np.int64 - ) + rng = np.random.default_rng(0) + values = rng.integers(low=-10000, high=10000, size=(4, 6)) data = da.from_array(values, chunks=(2, 2)) u = Variable(("x", "y"), values) v = Variable(("x", "y"), data) From e35c22fce1a72830acb5650871818294b60b3299 Mon Sep 17 00:00:00 2001 From: abrammer Date: Sat, 22 Apr 2023 13:16:59 -0600 Subject: [PATCH 08/17] Remove new type checking on test_variable edits Type checking throws errors on existing test lines --- xarray/tests/test_variable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 238cba86273..7fa9c344ba0 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -345,7 +345,7 @@ def test_pandas_period_index(self): assert "Period('2000-01-03', 'B')" in repr(v) @pytest.mark.parametrize("dtype", [float, int]) - def test_1d_math(self, dtype: np.typing.DTypeLike) -> None: + def test_1d_math(self, dtype): x = np.arange(5, dtype=dtype) y = np.ones(5, dtype=dtype) From 09033f54965f6dc44956352dc09dfb8eca3257f0 Mon Sep 17 00:00:00 2001 From: abrammer Date: Mon, 24 Apr 2023 10:38:54 -0600 Subject: [PATCH 09/17] Add typing to test_1d_math and ignore 1 existing line --- xarray/tests/test_variable.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 7fa9c344ba0..bef5efc15cc 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -345,7 +345,7 @@ def test_pandas_period_index(self): assert "Period('2000-01-03', 'B')" in repr(v) @pytest.mark.parametrize("dtype", [float, int]) - def test_1d_math(self, dtype): + def test_1d_math(self, dtype: np.typing.DTypeLike) -> None: x = np.arange(5, dtype=dtype) y = np.ones(5, dtype=dtype) @@ -367,7 +367,7 @@ def test_1d_math(self, dtype): assert_array_equal(v >> 2, x >> 2) # binary ops with numpy arrays assert_array_equal((v * x).values, x**2) - assert_array_equal((x * v).values, x**2) + assert_array_equal((x * v).values, x**2) # type: ignore[attr-defined] # TODO: Fix mypy thinking numpy takes priority, GH7780 assert_array_equal(v - y, v - 1) assert_array_equal(y - v, 1 - v) if dtype is int: From 6c26fd00c80d6af68f5dff12b24e0d0bd2204b0c Mon Sep 17 00:00:00 2001 From: abrammer Date: Mon, 24 Apr 2023 11:04:26 -0600 Subject: [PATCH 10/17] Add simple bitshift test on groupby ops --- xarray/tests/test_groupby.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 4488bda3eca..c195e7c7d44 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -808,6 +808,29 @@ def test_groupby_math_more() -> None: ds + ds.groupby("time.month") +def test_groupby_math_bitshift() -> None: + ds = Dataset( + { + "x": ("level", np.ones(4, dtype=int)), + "y": ("level", np.ones(4, dtype=int)*-1), + "level": [0, 1, 2, 3], + } + ) + left_expected = Dataset( + { + "x": ("level", [1, 2, 4, 8]), + "y": ("level", [-1, -2, -4, -8]), + "level": [0, 1, 2, 3] + } + ) + + left_actual = ds.groupby("level") << ds.level + assert_equal(left_expected, left_actual) + + right_actual = left_expected.groupby("level") >> left_expected.level + assert_equal(ds, right_actual) + + @pytest.mark.parametrize("use_flox", [True, False]) def test_groupby_bins_cut_kwargs(use_flox: bool) -> None: da = xr.DataArray(np.arange(12).reshape(6, 2), dims=("x", "y")) From f2837fe0805b8865cbc6839d15740b8654344599 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Apr 2023 17:06:07 +0000 Subject: [PATCH 11/17] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_groupby.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index c195e7c7d44..8cf36498c58 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -812,7 +812,7 @@ def test_groupby_math_bitshift() -> None: ds = Dataset( { "x": ("level", np.ones(4, dtype=int)), - "y": ("level", np.ones(4, dtype=int)*-1), + "y": ("level", np.ones(4, dtype=int) * -1), "level": [0, 1, 2, 3], } ) @@ -820,7 +820,7 @@ def test_groupby_math_bitshift() -> None: { "x": ("level", [1, 2, 4, 8]), "y": ("level", [-1, -2, -4, -8]), - "level": [0, 1, 2, 3] + "level": [0, 1, 2, 3], } ) From 453c3c130769564e947fc6770cda5712be5fd7d5 Mon Sep 17 00:00:00 2001 From: abrammer Date: Mon, 24 Apr 2023 12:39:20 -0600 Subject: [PATCH 12/17] Edit groupby bitshift test to use groups with len>1 --- xarray/tests/test_groupby.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 8cf36498c58..ca90b10af1f 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -809,25 +809,30 @@ def test_groupby_math_more() -> None: def test_groupby_math_bitshift() -> None: + # create new dataset of int's only ds = Dataset( { - "x": ("level", np.ones(4, dtype=int)), - "y": ("level", np.ones(4, dtype=int) * -1), - "level": [0, 1, 2, 3], + "x": ("index", np.ones(4, dtype=int)), + "y": ("index", np.ones(4, dtype=int)*-1), + "level":("index", [1, 1, 2, 2]), + "index": [0, 1, 2, 3] } ) + shift = DataArray([1, 2, 1], [("level", [1, 2, 8])]) + left_expected = Dataset( { - "x": ("level", [1, 2, 4, 8]), - "y": ("level", [-1, -2, -4, -8]), - "level": [0, 1, 2, 3], + "x": ("index", [2, 2, 4, 4]), + "y": ("index", [-2, -2, -4, -4]), + "level":("index", [2, 2, 8, 8]), + "index": [0, 1, 2, 3] } ) - left_actual = ds.groupby("level") << ds.level + left_actual = (ds.groupby("level") << shift).reset_coords(names='level') assert_equal(left_expected, left_actual) - right_actual = left_expected.groupby("level") >> left_expected.level + right_actual = (left_expected.groupby("level") >> shift).reset_coords(names='level') assert_equal(ds, right_actual) From e87533e6da4d23b2bcea35cf07be8bbbaa851f6e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Apr 2023 18:42:17 +0000 Subject: [PATCH 13/17] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_groupby.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index ca90b10af1f..48a9c17d2a8 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -813,9 +813,9 @@ def test_groupby_math_bitshift() -> None: ds = Dataset( { "x": ("index", np.ones(4, dtype=int)), - "y": ("index", np.ones(4, dtype=int)*-1), - "level":("index", [1, 1, 2, 2]), - "index": [0, 1, 2, 3] + "y": ("index", np.ones(4, dtype=int) * -1), + "level": ("index", [1, 1, 2, 2]), + "index": [0, 1, 2, 3], } ) shift = DataArray([1, 2, 1], [("level", [1, 2, 8])]) @@ -824,15 +824,15 @@ def test_groupby_math_bitshift() -> None: { "x": ("index", [2, 2, 4, 4]), "y": ("index", [-2, -2, -4, -4]), - "level":("index", [2, 2, 8, 8]), - "index": [0, 1, 2, 3] + "level": ("index", [2, 2, 8, 8]), + "index": [0, 1, 2, 3], } ) - left_actual = (ds.groupby("level") << shift).reset_coords(names='level') + left_actual = (ds.groupby("level") << shift).reset_coords(names="level") assert_equal(left_expected, left_actual) - right_actual = (left_expected.groupby("level") >> shift).reset_coords(names='level') + right_actual = (left_expected.groupby("level") >> shift).reset_coords(names="level") assert_equal(ds, right_actual) From 91633214d95e3d1e69f67ffee88c3e8ed11f39aa Mon Sep 17 00:00:00 2001 From: abrammer Date: Mon, 24 Apr 2023 14:35:02 -0600 Subject: [PATCH 14/17] Add rshift and lshift to docs Example use in computation.rst and entry in whats-new --- doc/user-guide/computation.rst | 4 ++++ doc/whats-new.rst | 2 ++ 2 files changed, 6 insertions(+) diff --git a/doc/user-guide/computation.rst b/doc/user-guide/computation.rst index d1f1274c7a1..69177aece29 100644 --- a/doc/user-guide/computation.rst +++ b/doc/user-guide/computation.rst @@ -63,6 +63,10 @@ Data arrays also implement many :py:class:`numpy.ndarray` methods: arr.round(2) arr.T + arr = xr.DataArray([0, 1, 2, 3, 4, 5]) + arr << 2 # only supported for int types + arr >> 1 + .. _missing_values: Missing values diff --git a/doc/whats-new.rst b/doc/whats-new.rst index bfc040eb271..bfee8bd33f0 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,6 +22,8 @@ v2023.05.0 (unreleased) New Features ~~~~~~~~~~~~ +- Add support for lshift and rshift binary operators (`<<`, `>>`) on :py:class:`xr.DataArray` of type :py:class:`int` (:issue:`7727` , :pull:`7741`) +By `Alan Brammer `_. Breaking changes From 46ccf4cfdc5f9800c566df2c60a935363c07d181 Mon Sep 17 00:00:00 2001 From: abrammer Date: Mon, 24 Apr 2023 17:04:35 -0600 Subject: [PATCH 15/17] Create new array in docs so examples later don't break --- doc/user-guide/computation.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/doc/user-guide/computation.rst b/doc/user-guide/computation.rst index 69177aece29..f8141f40321 100644 --- a/doc/user-guide/computation.rst +++ b/doc/user-guide/computation.rst @@ -63,9 +63,9 @@ Data arrays also implement many :py:class:`numpy.ndarray` methods: arr.round(2) arr.T - arr = xr.DataArray([0, 1, 2, 3, 4, 5]) - arr << 2 # only supported for int types - arr >> 1 + intarr = xr.DataArray([0, 1, 2, 3, 4, 5]) + intarr << 2 # only supported for int types + intarr >> 1 .. _missing_values: From 37bffa28abcfce0e5cfad9488388fc5a34245e79 Mon Sep 17 00:00:00 2001 From: abrammer Date: Tue, 25 Apr 2023 13:39:55 -0600 Subject: [PATCH 16/17] Indent second line on whats-new entry --- doc/whats-new.rst | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index bfee8bd33f0..cb69f764f2e 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,8 +22,9 @@ v2023.05.0 (unreleased) New Features ~~~~~~~~~~~~ -- Add support for lshift and rshift binary operators (`<<`, `>>`) on :py:class:`xr.DataArray` of type :py:class:`int` (:issue:`7727` , :pull:`7741`) -By `Alan Brammer `_. +- Add support for lshift and rshift binary operators (`<<`, `>>`) on + :py:class:`xr.DataArray` of type :py:class:`int` (:issue:`7727` , :pull:`7741`). + By `Alan Brammer `_. Breaking changes From 3bebcf8ece795bc2334b35b806a304d0fe34646a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 19:41:51 +0000 Subject: [PATCH 17/17] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index cb69f764f2e..f54e9ad9676 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,7 +22,7 @@ v2023.05.0 (unreleased) New Features ~~~~~~~~~~~~ -- Add support for lshift and rshift binary operators (`<<`, `>>`) on +- Add support for lshift and rshift binary operators (`<<`, `>>`) on :py:class:`xr.DataArray` of type :py:class:`int` (:issue:`7727` , :pull:`7741`). By `Alan Brammer `_.