Use set_options for asv bottleneck tests (#5986)

* Use set_options for bottleneck tests

* Use set_options in rolling

* Update rolling.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update rolling.py

* Update rolling.py

* set_options not needed.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Illviljan

asv_bench/asv.conf.json

@@ -62,7 +62,7 @@
         "pandas": [""],
         "netcdf4": [""],
         "scipy": [""],
-        "bottleneck": ["", null],
+        "bottleneck": [""],
         "dask": [""],
         "distributed": [""],
     },

asv_bench/benchmarks/dataarray_missing.py

@@ -16,13 +16,6 @@ def make_bench_data(shape, frac_nan, chunks):
     return da
 
 
-def requires_bottleneck():
-    try:
-        import bottleneck  # noqa: F401
-    except ImportError:
-        raise NotImplementedError()
-
-
 class DataArrayMissingInterpolateNA:
     def setup(self, shape, chunks, limit):
         if chunks is not None:
@@ -46,7 +39,6 @@ def time_interpolate_na(self, shape, chunks, limit):
 
 class DataArrayMissingBottleneck:
     def setup(self, shape, chunks, limit):
-        requires_bottleneck()
         if chunks is not None:
             requires_dask()
         self.da = make_bench_data(shape, 0.1, chunks)

asv_bench/benchmarks/rolling.py

@@ -36,29 +36,45 @@ def setup(self, *args, **kwargs):
             randn_long, dims="x", coords={"x": np.arange(long_nx) * 0.1}
         )
 
-    @parameterized(["func", "center"], (["mean", "count"], [True, False]))
-    def time_rolling(self, func, center):
-        getattr(self.ds.rolling(x=window, center=center), func)().load()
-
-    @parameterized(["func", "pandas"], (["mean", "count"], [True, False]))
-    def time_rolling_long(self, func, pandas):
+    @parameterized(
+        ["func", "center", "use_bottleneck"],
+        (["mean", "count"], [True, False], [True, False]),
+    )
+    def time_rolling(self, func, center, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            getattr(self.ds.rolling(x=window, center=center), func)().load()
+
+    @parameterized(
+        ["func", "pandas", "use_bottleneck"],
+        (["mean", "count"], [True, False], [True, False]),
+    )
+    def time_rolling_long(self, func, pandas, use_bottleneck):
         if pandas:
             se = self.da_long.to_series()
             getattr(se.rolling(window=window, min_periods=window), func)()
         else:
-            getattr(self.da_long.rolling(x=window, min_periods=window), func)().load()
-
-    @parameterized(["window_", "min_periods"], ([20, 40], [5, 5]))
-    def time_rolling_np(self, window_, min_periods):
-        self.ds.rolling(x=window_, center=False, min_periods=min_periods).reduce(
-            getattr(np, "nansum")
-        ).load()
-
-    @parameterized(["center", "stride"], ([True, False], [1, 1]))
-    def time_rolling_construct(self, center, stride):
-        self.ds.rolling(x=window, center=center).construct(
-            "window_dim", stride=stride
-        ).sum(dim="window_dim").load()
+            with xr.set_options(use_bottleneck=use_bottleneck):
+                getattr(
+                    self.da_long.rolling(x=window, min_periods=window), func
+                )().load()
+
+    @parameterized(
+        ["window_", "min_periods", "use_bottleneck"], ([20, 40], [5, 5], [True, False])
+    )
+    def time_rolling_np(self, window_, min_periods, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            self.ds.rolling(x=window_, center=False, min_periods=min_periods).reduce(
+                getattr(np, "nansum")
+            ).load()
+
+    @parameterized(
+        ["center", "stride", "use_bottleneck"], ([True, False], [1, 1], [True, False])
+    )
+    def time_rolling_construct(self, center, stride, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            self.ds.rolling(x=window, center=center).construct(
+                "window_dim", stride=stride
+            ).sum(dim="window_dim").load()
 
 
 class RollingDask(Rolling):
@@ -87,24 +103,28 @@ def setup(self, *args, **kwargs):
 
 
 class DataArrayRollingMemory(RollingMemory):
-    @parameterized("func", ["sum", "max", "mean"])
-    def peakmem_ndrolling_reduce(self, func):
-        roll = self.ds.var1.rolling(x=10, y=4)
-        getattr(roll, func)()
+    @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False]))
+    def peakmem_ndrolling_reduce(self, func, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            roll = self.ds.var1.rolling(x=10, y=4)
+            getattr(roll, func)()
 
-    @parameterized("func", ["sum", "max", "mean"])
-    def peakmem_1drolling_reduce(self, func):
-        roll = self.ds.var3.rolling(t=100)
-        getattr(roll, func)()
+    @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False]))
+    def peakmem_1drolling_reduce(self, func, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            roll = self.ds.var3.rolling(t=100)
+            getattr(roll, func)()
 
 
 class DatasetRollingMemory(RollingMemory):
-    @parameterized("func", ["sum", "max", "mean"])
-    def peakmem_ndrolling_reduce(self, func):
-        roll = self.ds.rolling(x=10, y=4)
-        getattr(roll, func)()
-
-    @parameterized("func", ["sum", "max", "mean"])
-    def peakmem_1drolling_reduce(self, func):
-        roll = self.ds.rolling(t=100)
-        getattr(roll, func)()
+    @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False]))
+    def peakmem_ndrolling_reduce(self, func, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            roll = self.ds.rolling(x=10, y=4)
+            getattr(roll, func)()
+
+    @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False]))
+    def peakmem_1drolling_reduce(self, func, use_bottleneck):
+        with xr.set_options(use_bottleneck=use_bottleneck):
+            roll = self.ds.rolling(t=100)
+            getattr(roll, func)()