Skip to content

Commit 9608307

Browse files
TomAugspurgerTomAugspurger
committed
TST: Replace yield-based tests in test_query_eval
1 parent b111691 commit 9608307

File tree

1 file changed

+50
-79
lines changed

1 file changed

+50
-79
lines changed

‎pandas/tests/frame/test_query_eval.py

+50-79
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
import operator
66
import pytest
7-
from itertools import product
87

98
from pandas.compat import (zip, range, lrange, StringIO)
109
from pandas import DataFrame, Series, Index, MultiIndex, date_range
@@ -165,7 +164,10 @@ def test_eval_resolvers_as_list(self):
165164

166165
class TestDataFrameQueryWithMultiIndex(tm.TestCase):
167166

168-
def check_query_with_named_multiindex(self, parser, engine):
167+
# parametrize single parser, for symmetry with other tests
168+
@pytest.mark.parametrize('parser', ['pandas'])
169+
@pytest.mark.parametrize('engine', ENGINES)
170+
def test_query_with_named_multiindex(self, parser, engine):
169171
tm.skip_if_no_ne(engine)
170172
a = np.random.choice(['red', 'green'], size=10)
171173
b = np.random.choice(['eggs', 'ham'], size=10)
@@ -214,11 +216,10 @@ def check_query_with_named_multiindex(self, parser, engine):
214216
assert_frame_equal(res1, exp)
215217
assert_frame_equal(res2, exp)
216218

217-
def test_query_with_named_multiindex(self):
218-
for parser, engine in product(['pandas'], ENGINES):
219-
yield self.check_query_with_named_multiindex, parser, engine
220-
221-
def check_query_with_unnamed_multiindex(self, parser, engine):
219+
# parametrize single parser, for symmetry with other tests
220+
@pytest.mark.parametrize('parser', ['pandas'])
221+
@pytest.mark.parametrize('engine', ENGINES)
222+
def test_query_with_unnamed_multiindex(self, parser, engine):
222223
tm.skip_if_no_ne(engine)
223224
a = np.random.choice(['red', 'green'], size=10)
224225
b = np.random.choice(['eggs', 'ham'], size=10)
@@ -308,11 +309,10 @@ def check_query_with_unnamed_multiindex(self, parser, engine):
308309
assert_frame_equal(res1, exp)
309310
assert_frame_equal(res2, exp)
310311

311-
def test_query_with_unnamed_multiindex(self):
312-
for parser, engine in product(['pandas'], ENGINES):
313-
yield self.check_query_with_unnamed_multiindex, parser, engine
314-
315-
def check_query_with_partially_named_multiindex(self, parser, engine):
312+
# parametrize single parser, for symmetry with other tests
313+
@pytest.mark.parametrize('parser', ['pandas'])
314+
@pytest.mark.parametrize('engine', ENGINES)
315+
def test_query_with_partially_named_multiindex(self, parser, engine):
316316
tm.skip_if_no_ne(engine)
317317
a = np.random.choice(['red', 'green'], size=10)
318318
b = np.arange(10)
@@ -341,17 +341,10 @@ def check_query_with_partially_named_multiindex(self, parser, engine):
341341
exp = df[ind != "red"]
342342
assert_frame_equal(res, exp)
343343

344-
def test_query_with_partially_named_multiindex(self):
345-
for parser, engine in product(['pandas'], ENGINES):
346-
yield (self.check_query_with_partially_named_multiindex,
347-
parser, engine)
348-
349-
def test_query_multiindex_get_index_resolvers(self):
350-
for parser, engine in product(['pandas'], ENGINES):
351-
yield (self.check_query_multiindex_get_index_resolvers, parser,
352-
engine)
353-
354-
def check_query_multiindex_get_index_resolvers(self, parser, engine):
344+
# parametrize single parser, for symmetry with other tests
345+
@pytest.mark.parametrize('parser', ['pandas'])
346+
@pytest.mark.parametrize('engine', ENGINES)
347+
def test_query_multiindex_get_index_resolvers(self, parser, engine):
355348
df = mkdf(10, 3, r_idx_nlevels=2, r_idx_names=['spam', 'eggs'])
356349
resolvers = df._get_index_resolvers()
357350

@@ -375,22 +368,18 @@ def to_series(mi, level):
375368
else:
376369
raise AssertionError("object must be a Series or Index")
377370

378-
def test_raise_on_panel_with_multiindex(self):
379-
for parser, engine in product(PARSERS, ENGINES):
380-
yield self.check_raise_on_panel_with_multiindex, parser, engine
381-
382-
def check_raise_on_panel_with_multiindex(self, parser, engine):
371+
@pytest.mark.parametrize('parser', PARSERS)
372+
@pytest.mark.parametrize('engine', ENGINES)
373+
def test_raise_on_panel_with_multiindex(self, parser, engine):
383374
tm.skip_if_no_ne()
384375
p = tm.makePanel(7)
385376
p.items = tm.makeCustomIndex(len(p.items), nlevels=2)
386377
with pytest.raises(NotImplementedError):
387378
pd.eval('p + 1', parser=parser, engine=engine)
388379

389-
def test_raise_on_panel4d_with_multiindex(self):
390-
for parser, engine in product(PARSERS, ENGINES):
391-
yield self.check_raise_on_panel4d_with_multiindex, parser, engine
392-
393-
def check_raise_on_panel4d_with_multiindex(self, parser, engine):
380+
@pytest.mark.parametrize('parser', PARSERS)
381+
@pytest.mark.parametrize('engine', ENGINES)
382+
def test_raise_on_panel4d_with_multiindex(self, parser, engine):
394383
tm.skip_if_no_ne()
395384
p4d = tm.makePanel4D(7)
396385
p4d.items = tm.makeCustomIndex(len(p4d.items), nlevels=2)
@@ -874,7 +863,9 @@ def test_query_builtin(self):
874863

875864
class TestDataFrameQueryStrings(tm.TestCase):
876865

877-
def check_str_query_method(self, parser, engine):
866+
@pytest.mark.parametrize('parser', PARSERS)
867+
@pytest.mark.parametrize('engine', ENGINES)
868+
def test_str_query_method(self, parser, engine):
878869
tm.skip_if_no_ne(engine)
879870
df = DataFrame(randn(10, 1), columns=['b'])
880871
df['strings'] = Series(list('aabbccddee'))
@@ -911,15 +902,9 @@ def check_str_query_method(self, parser, engine):
911902
assert_frame_equal(res, expect)
912903
assert_frame_equal(res, df[~df.strings.isin(['a'])])
913904

914-
def test_str_query_method(self):
915-
for parser, engine in product(PARSERS, ENGINES):
916-
yield self.check_str_query_method, parser, engine
917-
918-
def test_str_list_query_method(self):
919-
for parser, engine in product(PARSERS, ENGINES):
920-
yield self.check_str_list_query_method, parser, engine
921-
922-
def check_str_list_query_method(self, parser, engine):
905+
@pytest.mark.parametrize('parser', PARSERS)
906+
@pytest.mark.parametrize('engine', ENGINES)
907+
def test_str_list_query_method(self, parser, engine):
923908
tm.skip_if_no_ne(engine)
924909
df = DataFrame(randn(10, 1), columns=['b'])
925910
df['strings'] = Series(list('aabbccddee'))
@@ -958,7 +943,9 @@ def check_str_list_query_method(self, parser, engine):
958943
parser=parser)
959944
assert_frame_equal(res, expect)
960945

961-
def check_query_with_string_columns(self, parser, engine):
946+
@pytest.mark.parametrize('parser', PARSERS)
947+
@pytest.mark.parametrize('engine', ENGINES)
948+
def test_query_with_string_columns(self, parser, engine):
962949
tm.skip_if_no_ne(engine)
963950
df = DataFrame({'a': list('aaaabbbbcccc'),
964951
'b': list('aabbccddeeff'),
@@ -979,11 +966,9 @@ def check_query_with_string_columns(self, parser, engine):
979966
with pytest.raises(NotImplementedError):
980967
df.query('a in b and c < d', parser=parser, engine=engine)
981968

982-
def test_query_with_string_columns(self):
983-
for parser, engine in product(PARSERS, ENGINES):
984-
yield self.check_query_with_string_columns, parser, engine
985-
986-
def check_object_array_eq_ne(self, parser, engine):
969+
@pytest.mark.parametrize('parser', PARSERS)
970+
@pytest.mark.parametrize('engine', ENGINES)
971+
def test_object_array_eq_ne(self, parser, engine):
987972
tm.skip_if_no_ne(engine)
988973
df = DataFrame({'a': list('aaaabbbbcccc'),
989974
'b': list('aabbccddeeff'),
@@ -997,11 +982,9 @@ def check_object_array_eq_ne(self, parser, engine):
997982
exp = df[df.a != df.b]
998983
assert_frame_equal(res, exp)
999984

1000-
def test_object_array_eq_ne(self):
1001-
for parser, engine in product(PARSERS, ENGINES):
1002-
yield self.check_object_array_eq_ne, parser, engine
1003-
1004-
def check_query_with_nested_strings(self, parser, engine):
985+
@pytest.mark.parametrize('parser', PARSERS)
986+
@pytest.mark.parametrize('engine', ENGINES)
987+
def test_query_with_nested_strings(self, parser, engine):
1005988
tm.skip_if_no_ne(engine)
1006989
skip_if_no_pandas_parser(parser)
1007990
raw = """id event timestamp
@@ -1025,11 +1008,9 @@ def check_query_with_nested_strings(self, parser, engine):
10251008
engine=engine)
10261009
assert_frame_equal(expected, res)
10271010

1028-
def test_query_with_nested_string(self):
1029-
for parser, engine in product(PARSERS, ENGINES):
1030-
yield self.check_query_with_nested_strings, parser, engine
1031-
1032-
def check_query_with_nested_special_character(self, parser, engine):
1011+
@pytest.mark.parametrize('parser', PARSERS)
1012+
@pytest.mark.parametrize('engine', ENGINES)
1013+
def test_query_with_nested_special_character(self, parser, engine):
10331014
skip_if_no_pandas_parser(parser)
10341015
tm.skip_if_no_ne(engine)
10351016
df = DataFrame({'a': ['a', 'b', 'test & test'],
@@ -1038,12 +1019,9 @@ def check_query_with_nested_special_character(self, parser, engine):
10381019
expec = df[df.a == 'test & test']
10391020
assert_frame_equal(res, expec)
10401021

1041-
def test_query_with_nested_special_character(self):
1042-
for parser, engine in product(PARSERS, ENGINES):
1043-
yield (self.check_query_with_nested_special_character,
1044-
parser, engine)
1045-
1046-
def check_query_lex_compare_strings(self, parser, engine):
1022+
@pytest.mark.parametrize('parser', PARSERS)
1023+
@pytest.mark.parametrize('engine', ENGINES)
1024+
def test_query_lex_compare_strings(self, parser, engine):
10471025
tm.skip_if_no_ne(engine=engine)
10481026
import operator as opr
10491027

@@ -1058,11 +1036,9 @@ def check_query_lex_compare_strings(self, parser, engine):
10581036
expected = df[func(df.X, 'd')]
10591037
assert_frame_equal(res, expected)
10601038

1061-
def test_query_lex_compare_strings(self):
1062-
for parser, engine in product(PARSERS, ENGINES):
1063-
yield self.check_query_lex_compare_strings, parser, engine
1064-
1065-
def check_query_single_element_booleans(self, parser, engine):
1039+
@pytest.mark.parametrize('parser', PARSERS)
1040+
@pytest.mark.parametrize('engine', ENGINES)
1041+
def test_query_single_element_booleans(self, parser, engine):
10661042
tm.skip_if_no_ne(engine)
10671043
columns = 'bid', 'bidsize', 'ask', 'asksize'
10681044
data = np.random.randint(2, size=(1, len(columns))).astype(bool)
@@ -1071,11 +1047,10 @@ def check_query_single_element_booleans(self, parser, engine):
10711047
expected = df[df.bid & df.ask]
10721048
assert_frame_equal(res, expected)
10731049

1074-
def test_query_single_element_booleans(self):
1075-
for parser, engine in product(PARSERS, ENGINES):
1076-
yield self.check_query_single_element_booleans, parser, engine
1077-
1078-
def check_query_string_scalar_variable(self, parser, engine):
1050+
# parametrize single parser, for symmetry with other tests
1051+
@pytest.mark.parametrize('parser', ['pandas'])
1052+
@pytest.mark.parametrize('engine', ENGINES)
1053+
def test_query_string_scalar_variable(self, parser, engine):
10791054
tm.skip_if_no_ne(engine)
10801055
df = pd.DataFrame({'Symbol': ['BUD US', 'BUD US', 'IBM US', 'IBM US'],
10811056
'Price': [109.70, 109.72, 183.30, 183.35]})
@@ -1084,10 +1059,6 @@ def check_query_string_scalar_variable(self, parser, engine):
10841059
r = df.query('Symbol == @symb', parser=parser, engine=engine)
10851060
assert_frame_equal(e, r)
10861061

1087-
def test_query_string_scalar_variable(self):
1088-
for parser, engine in product(['pandas'], ENGINES):
1089-
yield self.check_query_string_scalar_variable, parser, engine
1090-
10911062

10921063
class TestDataFrameEvalNumExprPandas(tm.TestCase):
10931064

0 commit comments

Comments
 (0)