TST: make tests involving use_numexpr more robust (#44601)

Author: jorisvandenbossche

pandas/_testing/contexts.py

@@ -14,6 +14,8 @@
 
 import numpy as np
 
+from pandas import set_option
+
 from pandas.io.common import get_handle
 
 
@@ -202,11 +204,11 @@ def use_numexpr(use, min_elements=None):
 
     olduse = expr.USE_NUMEXPR
     oldmin = expr._MIN_ELEMENTS
-    expr.set_use_numexpr(use)
+    set_option("compute.use_numexpr", use)
     expr._MIN_ELEMENTS = min_elements
     yield
     expr._MIN_ELEMENTS = oldmin
-    expr.set_use_numexpr(olduse)
+    set_option("compute.use_numexpr", olduse)
 
 
 class RNGContext:

pandas/tests/test_downstream.py

@@ -10,6 +10,7 @@
 
 import pandas.util._test_decorators as td
 
+import pandas as pd
 from pandas import DataFrame
 import pandas._testing as tm
 
@@ -32,14 +33,21 @@ def df():
 @pytest.mark.filterwarnings("ignore:.*64Index is deprecated:FutureWarning")
 def test_dask(df):
 
-    toolz = import_module("toolz")  # noqa:F841
-    dask = import_module("dask")  # noqa:F841
+    # dask sets "compute.use_numexpr" to False, so catch the current value
+    # and ensure to reset it afterwards to avoid impacting other tests
+    olduse = pd.get_option("compute.use_numexpr")
 
-    import dask.dataframe as dd
+    try:
+        toolz = import_module("toolz")  # noqa:F841
+        dask = import_module("dask")  # noqa:F841
+
+        import dask.dataframe as dd
 
-    ddf = dd.from_pandas(df, npartitions=3)
-    assert ddf.A is not None
-    assert ddf.compute() is not None
+        ddf = dd.from_pandas(df, npartitions=3)
+        assert ddf.A is not None
+        assert ddf.compute() is not None
+    finally:
+        pd.set_option("compute.use_numexpr", olduse)
 
 
 def test_xarray(df):

pandas/tests/test_expressions.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 
+from pandas import set_option
 import pandas._testing as tm
 from pandas.core.api import (
     DataFrame,
@@ -65,9 +66,9 @@ def call_op(df, other, flex: bool, opname: str):
         else:
             op = getattr(operator, opname)
 
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         expected = op(df, other)
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
 
         expr.get_test_result()
 
@@ -107,9 +108,9 @@ def run_binary(self, df, other, flex: bool):
     def run_frame(self, df, other, flex: bool):
         self.run_arithmetic(df, other, flex)
 
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         binary_comp = other + 1
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
         self.run_binary(df, binary_comp, flex)
 
         for i in range(len(df.columns)):
@@ -179,9 +180,9 @@ def testit():
             result = expr._can_use_numexpr(op, op_str, right, right, "evaluate")
             assert not result
 
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         testit()
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
         expr.set_numexpr_threads(1)
         testit()
         expr.set_numexpr_threads()
@@ -215,9 +216,9 @@ def testit():
             result = expr._can_use_numexpr(op, op_str, right, f22, "evaluate")
             assert not result
 
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         testit()
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
         expr.set_numexpr_threads(1)
         testit()
         expr.set_numexpr_threads()
@@ -233,9 +234,9 @@ def testit():
             expected = np.where(c, df.values, df.values + 1)
             tm.assert_numpy_array_equal(result, expected)
 
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         testit()
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
         expr.set_numexpr_threads(1)
         testit()
         expr.set_numexpr_threads()
@@ -360,9 +361,9 @@ def test_frame_series_axis(self, axis, arith):
 
         op_func = getattr(df, arith)
 
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         expected = op_func(other, axis=axis)
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
 
         result = op_func(other, axis=axis)
         tm.assert_frame_equal(expected, result)
@@ -387,9 +388,9 @@ def test_python_semantics_with_numexpr_installed(self, op, box, scalar):
         result = method(scalar)
 
         # compare result with numpy
-        expr.set_use_numexpr(False)
+        set_option("compute.use_numexpr", False)
         expected = method(scalar)
-        expr.set_use_numexpr(True)
+        set_option("compute.use_numexpr", True)
         tm.assert_equal(result, expected)
 
         # compare result element-wise with Python