Skip to content

Commit da22c30

Browse files
SaturnFromTitanproost
SaturnFromTitan
authored and
proost
committed
Fixed test util imports in pandas/tests/indexing (pandas-dev#29291)
xref: pandas-dev#29272
1 parent 0a590dd commit da22c30

File tree

6 files changed

+82
-88
lines changed

6 files changed

+82
-88
lines changed

‎pandas/tests/indexing/test_categorical.py

+42-43
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,7 @@
1515
Timestamp,
1616
)
1717
from pandas.api.types import CategoricalDtype as CDT
18-
from pandas.util import testing as tm
19-
from pandas.util.testing import assert_frame_equal, assert_series_equal
18+
import pandas.util.testing as tm
2019

2120

2221
class TestCategoricalIndex:
@@ -52,7 +51,7 @@ def test_loc_scalar(self):
5251
expected = DataFrame(
5352
{"A": [0, 1, 5], "B": (Series(list("aaa")).astype(CDT(list("cab"))))}
5453
).set_index("B")
55-
assert_frame_equal(result, expected)
54+
tm.assert_frame_equal(result, expected)
5655

5756
df = self.df.copy()
5857
df.loc["a"] = 20
@@ -62,7 +61,7 @@ def test_loc_scalar(self):
6261
"B": (Series(list("aabbca")).astype(CDT(list("cab")))),
6362
}
6463
).set_index("B")
65-
assert_frame_equal(df, expected)
64+
tm.assert_frame_equal(df, expected)
6665

6766
# value not in the categories
6867
with pytest.raises(KeyError, match=r"^'d'$"):
@@ -329,12 +328,12 @@ def test_loc_listlike(self):
329328
# list of labels
330329
result = self.df.loc[["c", "a"]]
331330
expected = self.df.iloc[[4, 0, 1, 5]]
332-
assert_frame_equal(result, expected, check_index_type=True)
331+
tm.assert_frame_equal(result, expected, check_index_type=True)
333332

334333
result = self.df2.loc[["a", "b", "e"]]
335334
exp_index = CategoricalIndex(list("aaabbe"), categories=list("cabe"), name="B")
336335
expected = DataFrame({"A": [0, 1, 5, 2, 3, np.nan]}, index=exp_index)
337-
assert_frame_equal(result, expected, check_index_type=True)
336+
tm.assert_frame_equal(result, expected, check_index_type=True)
338337

339338
# element in the categories but not in the values
340339
with pytest.raises(KeyError, match=r"^'e'$"):
@@ -346,13 +345,13 @@ def test_loc_listlike(self):
346345
result = df.loc[["a", "b", "e"]]
347346
exp_index = CategoricalIndex(list("aaabbe"), categories=list("cabe"), name="B")
348347
expected = DataFrame({"A": [0, 1, 5, 2, 3, 20]}, index=exp_index)
349-
assert_frame_equal(result, expected)
348+
tm.assert_frame_equal(result, expected)
350349

351350
df = self.df2.copy()
352351
result = df.loc[["a", "b", "e"]]
353352
exp_index = CategoricalIndex(list("aaabbe"), categories=list("cabe"), name="B")
354353
expected = DataFrame({"A": [0, 1, 5, 2, 3, np.nan]}, index=exp_index)
355-
assert_frame_equal(result, expected, check_index_type=True)
354+
tm.assert_frame_equal(result, expected, check_index_type=True)
356355

357356
# not all labels in the categories
358357
with pytest.raises(
@@ -474,7 +473,7 @@ def test_getitem_with_listlike(self):
474473
)
475474
dummies = pd.get_dummies(cats)
476475
result = dummies[[c for c in dummies.columns]]
477-
assert_frame_equal(result, expected)
476+
tm.assert_frame_equal(result, expected)
478477

479478
def test_setitem_listlike(self):
480479

@@ -500,18 +499,18 @@ def test_ix_categorical_index(self):
500499
cdf.columns = CategoricalIndex(df.columns)
501500

502501
expect = Series(df.loc["A", :], index=cdf.columns, name="A")
503-
assert_series_equal(cdf.loc["A", :], expect)
502+
tm.assert_series_equal(cdf.loc["A", :], expect)
504503

505504
expect = Series(df.loc[:, "X"], index=cdf.index, name="X")
506-
assert_series_equal(cdf.loc[:, "X"], expect)
505+
tm.assert_series_equal(cdf.loc[:, "X"], expect)
507506

508507
exp_index = CategoricalIndex(list("AB"), categories=["A", "B", "C"])
509508
expect = DataFrame(df.loc[["A", "B"], :], columns=cdf.columns, index=exp_index)
510-
assert_frame_equal(cdf.loc[["A", "B"], :], expect)
509+
tm.assert_frame_equal(cdf.loc[["A", "B"], :], expect)
511510

512511
exp_columns = CategoricalIndex(list("XY"), categories=["X", "Y", "Z"])
513512
expect = DataFrame(df.loc[:, ["X", "Y"]], index=cdf.index, columns=exp_columns)
514-
assert_frame_equal(cdf.loc[:, ["X", "Y"]], expect)
513+
tm.assert_frame_equal(cdf.loc[:, ["X", "Y"]], expect)
515514

516515
# non-unique
517516
df = DataFrame(np.random.randn(3, 3), index=list("ABA"), columns=list("XYX"))
@@ -521,25 +520,25 @@ def test_ix_categorical_index(self):
521520

522521
exp_index = CategoricalIndex(list("AA"), categories=["A", "B"])
523522
expect = DataFrame(df.loc["A", :], columns=cdf.columns, index=exp_index)
524-
assert_frame_equal(cdf.loc["A", :], expect)
523+
tm.assert_frame_equal(cdf.loc["A", :], expect)
525524

526525
exp_columns = CategoricalIndex(list("XX"), categories=["X", "Y"])
527526
expect = DataFrame(df.loc[:, "X"], index=cdf.index, columns=exp_columns)
528-
assert_frame_equal(cdf.loc[:, "X"], expect)
527+
tm.assert_frame_equal(cdf.loc[:, "X"], expect)
529528

530529
expect = DataFrame(
531530
df.loc[["A", "B"], :],
532531
columns=cdf.columns,
533532
index=CategoricalIndex(list("AAB")),
534533
)
535-
assert_frame_equal(cdf.loc[["A", "B"], :], expect)
534+
tm.assert_frame_equal(cdf.loc[["A", "B"], :], expect)
536535

537536
expect = DataFrame(
538537
df.loc[:, ["X", "Y"]],
539538
index=cdf.index,
540539
columns=CategoricalIndex(list("XXY")),
541540
)
542-
assert_frame_equal(cdf.loc[:, ["X", "Y"]], expect)
541+
tm.assert_frame_equal(cdf.loc[:, ["X", "Y"]], expect)
543542

544543
def test_read_only_source(self):
545544
# GH 10043
@@ -550,15 +549,15 @@ def test_read_only_source(self):
550549
ro_array.setflags(write=False)
551550
ro_df = DataFrame(ro_array)
552551

553-
assert_frame_equal(rw_df.iloc[[1, 2, 3]], ro_df.iloc[[1, 2, 3]])
554-
assert_frame_equal(rw_df.iloc[[1]], ro_df.iloc[[1]])
555-
assert_series_equal(rw_df.iloc[1], ro_df.iloc[1])
556-
assert_frame_equal(rw_df.iloc[1:3], ro_df.iloc[1:3])
552+
tm.assert_frame_equal(rw_df.iloc[[1, 2, 3]], ro_df.iloc[[1, 2, 3]])
553+
tm.assert_frame_equal(rw_df.iloc[[1]], ro_df.iloc[[1]])
554+
tm.assert_series_equal(rw_df.iloc[1], ro_df.iloc[1])
555+
tm.assert_frame_equal(rw_df.iloc[1:3], ro_df.iloc[1:3])
557556

558-
assert_frame_equal(rw_df.loc[[1, 2, 3]], ro_df.loc[[1, 2, 3]])
559-
assert_frame_equal(rw_df.loc[[1]], ro_df.loc[[1]])
560-
assert_series_equal(rw_df.loc[1], ro_df.loc[1])
561-
assert_frame_equal(rw_df.loc[1:3], ro_df.loc[1:3])
557+
tm.assert_frame_equal(rw_df.loc[[1, 2, 3]], ro_df.loc[[1, 2, 3]])
558+
tm.assert_frame_equal(rw_df.loc[[1]], ro_df.loc[[1]])
559+
tm.assert_series_equal(rw_df.loc[1], ro_df.loc[1])
560+
tm.assert_frame_equal(rw_df.loc[1:3], ro_df.loc[1:3])
562561

563562
def test_reindexing(self):
564563
df = DataFrame(
@@ -574,19 +573,19 @@ def test_reindexing(self):
574573
expected = DataFrame({"A": [0, 1, np.nan], "B": Series(list("abe"))}).set_index(
575574
"B"
576575
)
577-
assert_frame_equal(result, expected, check_index_type=True)
576+
tm.assert_frame_equal(result, expected, check_index_type=True)
578577

579578
result = df.reindex(["a", "b"])
580579
expected = DataFrame({"A": [0, 1], "B": Series(list("ab"))}).set_index("B")
581-
assert_frame_equal(result, expected, check_index_type=True)
580+
tm.assert_frame_equal(result, expected, check_index_type=True)
582581

583582
result = df.reindex(["e"])
584583
expected = DataFrame({"A": [np.nan], "B": Series(["e"])}).set_index("B")
585-
assert_frame_equal(result, expected, check_index_type=True)
584+
tm.assert_frame_equal(result, expected, check_index_type=True)
586585

587586
result = df.reindex(["d"])
588587
expected = DataFrame({"A": [np.nan], "B": Series(["d"])}).set_index("B")
589-
assert_frame_equal(result, expected, check_index_type=True)
588+
tm.assert_frame_equal(result, expected, check_index_type=True)
590589

591590
# since we are actually reindexing with a Categorical
592591
# then return a Categorical
@@ -596,40 +595,40 @@ def test_reindexing(self):
596595
expected = DataFrame(
597596
{"A": [0, np.nan], "B": Series(list("ae")).astype(CDT(cats))}
598597
).set_index("B")
599-
assert_frame_equal(result, expected, check_index_type=True)
598+
tm.assert_frame_equal(result, expected, check_index_type=True)
600599

601600
result = df.reindex(Categorical(["a"], categories=cats))
602601
expected = DataFrame(
603602
{"A": [0], "B": Series(list("a")).astype(CDT(cats))}
604603
).set_index("B")
605-
assert_frame_equal(result, expected, check_index_type=True)
604+
tm.assert_frame_equal(result, expected, check_index_type=True)
606605

607606
result = df.reindex(["a", "b", "e"])
608607
expected = DataFrame({"A": [0, 1, np.nan], "B": Series(list("abe"))}).set_index(
609608
"B"
610609
)
611-
assert_frame_equal(result, expected, check_index_type=True)
610+
tm.assert_frame_equal(result, expected, check_index_type=True)
612611

613612
result = df.reindex(["a", "b"])
614613
expected = DataFrame({"A": [0, 1], "B": Series(list("ab"))}).set_index("B")
615-
assert_frame_equal(result, expected, check_index_type=True)
614+
tm.assert_frame_equal(result, expected, check_index_type=True)
616615

617616
result = df.reindex(["e"])
618617
expected = DataFrame({"A": [np.nan], "B": Series(["e"])}).set_index("B")
619-
assert_frame_equal(result, expected, check_index_type=True)
618+
tm.assert_frame_equal(result, expected, check_index_type=True)
620619

621620
# give back the type of categorical that we received
622621
result = df.reindex(Categorical(["a", "e"], categories=cats, ordered=True))
623622
expected = DataFrame(
624623
{"A": [0, np.nan], "B": Series(list("ae")).astype(CDT(cats, ordered=True))}
625624
).set_index("B")
626-
assert_frame_equal(result, expected, check_index_type=True)
625+
tm.assert_frame_equal(result, expected, check_index_type=True)
627626

628627
result = df.reindex(Categorical(["a", "d"], categories=["a", "d"]))
629628
expected = DataFrame(
630629
{"A": [0, np.nan], "B": Series(list("ad")).astype(CDT(["a", "d"]))}
631630
).set_index("B")
632-
assert_frame_equal(result, expected, check_index_type=True)
631+
tm.assert_frame_equal(result, expected, check_index_type=True)
633632

634633
# passed duplicate indexers are not allowed
635634
msg = "cannot reindex from a duplicate axis"
@@ -661,7 +660,7 @@ def test_loc_slice(self):
661660

662661
# result = df.loc[1:5]
663662
# expected = df.iloc[[1,2,3,4]]
664-
# assert_frame_equal(result, expected)
663+
# tm.assert_frame_equal(result, expected)
665664

666665
def test_loc_and_at_with_categorical_index(self):
667666
# GH 20629
@@ -681,19 +680,19 @@ def test_boolean_selection(self):
681680

682681
result = df3[df3.index == "a"]
683682
expected = df3.iloc[[]]
684-
assert_frame_equal(result, expected)
683+
tm.assert_frame_equal(result, expected)
685684

686685
result = df4[df4.index == "a"]
687686
expected = df4.iloc[[]]
688-
assert_frame_equal(result, expected)
687+
tm.assert_frame_equal(result, expected)
689688

690689
result = df3[df3.index == 1]
691690
expected = df3.iloc[[0, 1, 3]]
692-
assert_frame_equal(result, expected)
691+
tm.assert_frame_equal(result, expected)
693692

694693
result = df4[df4.index == 1]
695694
expected = df4.iloc[[0, 1, 3]]
696-
assert_frame_equal(result, expected)
695+
tm.assert_frame_equal(result, expected)
697696

698697
# since we have an ordered categorical
699698

@@ -703,11 +702,11 @@ def test_boolean_selection(self):
703702
# name='B')
704703
result = df3[df3.index < 2]
705704
expected = df3.iloc[[4]]
706-
assert_frame_equal(result, expected)
705+
tm.assert_frame_equal(result, expected)
707706

708707
result = df3[df3.index > 1]
709708
expected = df3.iloc[[]]
710-
assert_frame_equal(result, expected)
709+
tm.assert_frame_equal(result, expected)
711710

712711
# unordered
713712
# cannot be compared

0 commit comments

Comments
 (0)