diff --git a/python/cudf/cudf/tests/test_dataframe.py b/python/cudf/cudf/tests/test_dataframe.py index ba2caf7c6c8..114d54ceb86 100644 --- a/python/cudf/cudf/tests/test_dataframe.py +++ b/python/cudf/cudf/tests/test_dataframe.py @@ -246,17 +246,15 @@ def test_series_init_none(): sr1 = cudf.Series() got = sr1.to_string() - expect = sr1.to_pandas().__repr__() - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = repr(sr1.to_pandas()) + assert got == expect # 2: Using `None` as an initializer sr2 = cudf.Series(None) got = sr2.to_string() - expect = sr2.to_pandas().__repr__() - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = repr(sr2.to_pandas()) + assert got == expect def test_dataframe_basic(): @@ -843,21 +841,20 @@ def test_dataframe_to_string_with_masked_data(): def test_dataframe_to_string_wide(monkeypatch): monkeypatch.setenv("COLUMNS", "79") # Test basic - df = cudf.DataFrame() - for i in range(100): - df["a{}".format(i)] = list(range(3)) - pd.options.display.max_columns = 0 - got = df.to_string() + df = cudf.DataFrame({f"a{i}": [0, 1, 2] for i in range(100)}) + with pd.option_context("display.max_columns", 0): + got = df.to_string() - expect = """ - a0 a1 a2 a3 a4 a5 a6 a7 ... a92 a93 a94 a95 a96 a97 a98 a99 -0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 -1 1 1 1 1 1 1 1 1 ... 1 1 1 1 1 1 1 1 -2 2 2 2 2 2 2 2 2 ... 2 2 2 2 2 2 2 2 -[3 rows x 100 columns] -""" - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = textwrap.dedent( + """\ + a0 a1 a2 a3 a4 a5 a6 a7 ... a92 a93 a94 a95 a96 a97 a98 a99 + 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 + 1 1 1 1 1 1 1 1 1 ... 1 1 1 1 1 1 1 1 + 2 2 2 2 2 2 2 2 2 ... 2 2 2 2 2 2 2 2 + + [3 rows x 100 columns]""" # noqa: E501 + ) + assert got == expect def test_dataframe_empty_to_string(): @@ -865,9 +862,8 @@ def test_dataframe_empty_to_string(): df = cudf.DataFrame() got = df.to_string() - expect = "Empty DataFrame\nColumns: []\nIndex: []\n" - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = "Empty DataFrame\nColumns: []\nIndex: []" + assert got == expect def test_dataframe_emptycolumns_to_string(): @@ -877,9 +873,8 @@ def test_dataframe_emptycolumns_to_string(): df["b"] = [] got = df.to_string() - expect = "Empty DataFrame\nColumns: [a, b]\nIndex: []\n" - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = "Empty DataFrame\nColumns: [a, b]\nIndex: []" + assert got == expect def test_dataframe_copy(): @@ -890,14 +885,14 @@ def test_dataframe_copy(): df2["b"] = [4, 5, 6] got = df.to_string() - expect = """ - a -0 1 -1 2 -2 3 -""" - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = textwrap.dedent( + """\ + a + 0 1 + 1 2 + 2 3""" + ) + assert got == expect def test_dataframe_copy_shallow(): @@ -908,14 +903,14 @@ def test_dataframe_copy_shallow(): df2["b"] = [4, 2, 3] got = df.to_string() - expect = """ - a -0 1 -1 2 -2 3 -""" - # values should match despite whitespace difference - assert got.split() == expect.split() + expect = textwrap.dedent( + """\ + a + 0 1 + 1 2 + 2 3""" + ) + assert got == expect def test_dataframe_dtypes():