def test_sort_1(self):
df = create_df()
df["U"] = [9, 8, 7, 1, 2]
w = Wrap(df)
wnew = w.sort("U + 1")
self.assertListEqual([0, 1, 2, 4, 3], list(wnew.df["B"]))
self.assertListEqual(list(range(5)), list(w.df["B"]))
def test_sort_asc_desc(self):
df = create_df()
df["U"] = [9, 9, 7, 1, 2]
df["U2"] = [1, 2, 3, 4, 5]
w = Wrap(df)
wnew = w.sort("U asc, U2 desc")
assert [3, 4, 2, 1, 0] == list(wnew.df["B"])
assert [1, 2, 7, 9, 9] == list(wnew.df["U"])
assert [4, 5, 3, 2, 1] == list(wnew.df["U2"])
def test_top_asc(self):
df = create_df()
df["U"] = [9, 8, 7, 1, 2]
w = Wrap(df)
wnew = w.top(4, "U + 1 asc")
assert [1, 2, 7, 8] == list(wnew.df["U"])
assert 6 == len(wnew.df.columns)
wexpected = w.sort("U asc").take(4)
pd.testing.assert_frame_equal(wnew.df, wexpected.df)
def test_sort(self):
df = create_df()
df["U"] = [9, 1, 7, 1, 2]
expected = [1, 3, 4, 2, 0]
w = Wrap(df)
wnew = w.sort(["U + 1 asc", "B asc"])
self.assertListEqual([1, 1, 2, 7, 9], list(wnew.df["U"]))
self.assertListEqual(expected, list(wnew.df["B"]))
self.assertListEqual(list(range(5)), list(w.df["B"]))
assert 6 == len(wnew.df.columns)
def test_top(self):
df = create_df()
df["U"] = [9, 8, 7, 1, 2]
w = Wrap(df)
wnew = w.top(4, "U + 1")
assert [9, 8, 7, 2] == list(wnew.df["U"])
self.assertListEqual([0, 1, 2, 4], list(wnew.df["B"]))
self.assertListEqual(["foo1", "foo2", "foo3", "foo5"],
list(wnew.df["C"]))
assert 6 == len(wnew.df.columns)
wexpected = w.sort("U").take(4)
pd.testing.assert_frame_equal(wnew.df, wexpected.df)
