TST/REF: io/parser/(test_dtypes.py, test_usecols.py)

pandas/tests/io/parser/dtypes/test_categorical.py

@@ -0,0 +1,294 @@
+"""
+Tests dtype specification during parsing
+for all of the parsers defined in parsers.py
+"""
+from io import StringIO
+import os
+
+import numpy as np
+import pytest
+
+from pandas.core.dtypes.dtypes import CategoricalDtype
+
+import pandas as pd
+from pandas import Categorical, DataFrame, Timestamp
+import pandas._testing as tm
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        "category",
+        CategoricalDtype(),
+        {"a": "category", "b": "category", "c": CategoricalDtype()},
+    ],
+)
+def test_categorical_dtype(all_parsers, dtype):
+    # see gh-10153
+    parser = all_parsers
+    data = """a,b,c
+1,a,3.4
+1,a,3.4
+2,b,4.5"""
+    expected = DataFrame(
+        {
+            "a": Categorical(["1", "1", "2"]),
+            "b": Categorical(["a", "a", "b"]),
+            "c": Categorical(["3.4", "3.4", "4.5"]),
+        }
+    )
+    actual = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(actual, expected)
+
+
+@pytest.mark.parametrize("dtype", [{"b": "category"}, {1: "category"}])
+def test_categorical_dtype_single(all_parsers, dtype):
+    # see gh-10153
+    parser = all_parsers
+    data = """a,b,c
+1,a,3.4
+1,a,3.4
+2,b,4.5"""
+    expected = DataFrame(
+        {"a": [1, 1, 2], "b": Categorical(["a", "a", "b"]), "c": [3.4, 3.4, 4.5]}
+    )
+    actual = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(actual, expected)
+
+
+def test_categorical_dtype_unsorted(all_parsers):
+    # see gh-10153
+    parser = all_parsers
+    data = """a,b,c
+1,b,3.4
+1,b,3.4
+2,a,4.5"""
+    expected = DataFrame(
+        {
+            "a": Categorical(["1", "1", "2"]),
+            "b": Categorical(["b", "b", "a"]),
+            "c": Categorical(["3.4", "3.4", "4.5"]),
+        }
+    )
+    actual = parser.read_csv(StringIO(data), dtype="category")
+    tm.assert_frame_equal(actual, expected)
+
+
+def test_categorical_dtype_missing(all_parsers):
+    # see gh-10153
+    parser = all_parsers
+    data = """a,b,c
+1,b,3.4
+1,nan,3.4
+2,a,4.5"""
+    expected = DataFrame(
+        {
+            "a": Categorical(["1", "1", "2"]),
+            "b": Categorical(["b", np.nan, "a"]),
+            "c": Categorical(["3.4", "3.4", "4.5"]),
+        }
+    )
+    actual = parser.read_csv(StringIO(data), dtype="category")
+    tm.assert_frame_equal(actual, expected)
+
+
+@pytest.mark.slow
+def test_categorical_dtype_high_cardinality_numeric(all_parsers):
+    # see gh-18186
+    parser = all_parsers
+    data = np.sort([str(i) for i in range(524289)])
+    expected = DataFrame({"a": Categorical(data, ordered=True)})
+
+    actual = parser.read_csv(StringIO("a\n" + "\n".join(data)), dtype="category")
+    actual["a"] = actual["a"].cat.reorder_categories(
+        np.sort(actual.a.cat.categories), ordered=True
+    )
+    tm.assert_frame_equal(actual, expected)
+
+
+def test_categorical_dtype_utf16(all_parsers, csv_dir_path):
+    # see gh-10153
+    pth = os.path.join(csv_dir_path, "utf16_ex.txt")
+    parser = all_parsers
+    encoding = "utf-16"
+    sep = "\t"
+
+    expected = parser.read_csv(pth, sep=sep, encoding=encoding)
+    expected = expected.apply(Categorical)
+
+    actual = parser.read_csv(pth, sep=sep, encoding=encoding, dtype="category")
+    tm.assert_frame_equal(actual, expected)
+
+
+def test_categorical_dtype_chunksize_infer_categories(all_parsers):
+    # see gh-10153
+    parser = all_parsers
+    data = """a,b
+1,a
+1,b
+1,b
+2,c"""
+    expecteds = [
+        DataFrame({"a": [1, 1], "b": Categorical(["a", "b"])}),
+        DataFrame({"a": [1, 2], "b": Categorical(["b", "c"])}, index=[2, 3]),
+    ]
+    with parser.read_csv(
+        StringIO(data), dtype={"b": "category"}, chunksize=2
+    ) as actuals:
+        for actual, expected in zip(actuals, expecteds):
+            tm.assert_frame_equal(actual, expected)
+
+
+def test_categorical_dtype_chunksize_explicit_categories(all_parsers):
+    # see gh-10153
+    parser = all_parsers
+    data = """a,b
+1,a
+1,b
+1,b
+2,c"""
+    cats = ["a", "b", "c"]
+    expecteds = [
+        DataFrame({"a": [1, 1], "b": Categorical(["a", "b"], categories=cats)}),
+        DataFrame(
+            {"a": [1, 2], "b": Categorical(["b", "c"], categories=cats)},
+            index=[2, 3],
+        ),
+    ]
+    dtype = CategoricalDtype(cats)
+    with parser.read_csv(StringIO(data), dtype={"b": dtype}, chunksize=2) as actuals:
+        for actual, expected in zip(actuals, expecteds):
+            tm.assert_frame_equal(actual, expected)
+
+
+def test_categorical_dtype_latin1(all_parsers, csv_dir_path):
+    # see gh-10153
+    pth = os.path.join(csv_dir_path, "unicode_series.csv")
+    parser = all_parsers
+    encoding = "latin-1"
+
+    expected = parser.read_csv(pth, header=None, encoding=encoding)
+    expected[1] = Categorical(expected[1])
+
+    actual = parser.read_csv(pth, header=None, encoding=encoding, dtype={1: "category"})
+    tm.assert_frame_equal(actual, expected)
+
+
+@pytest.mark.parametrize("ordered", [False, True])
+@pytest.mark.parametrize(
+    "categories",
+    [["a", "b", "c"], ["a", "c", "b"], ["a", "b", "c", "d"], ["c", "b", "a"]],
+)
+def test_categorical_category_dtype(all_parsers, categories, ordered):
+    parser = all_parsers
+    data = """a,b
+1,a
+1,b
+1,b
+2,c"""
+    expected = DataFrame(
+        {
+            "a": [1, 1, 1, 2],
+            "b": Categorical(
+                ["a", "b", "b", "c"], categories=categories, ordered=ordered
+            ),
+        }
+    )
+
+    dtype = {"b": CategoricalDtype(categories=categories, ordered=ordered)}
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)
+
+
+def test_categorical_category_dtype_unsorted(all_parsers):
+    parser = all_parsers
+    data = """a,b
+1,a
+1,b
+1,b
+2,c"""
+    dtype = CategoricalDtype(["c", "b", "a"])
+    expected = DataFrame(
+        {
+            "a": [1, 1, 1, 2],
+            "b": Categorical(["a", "b", "b", "c"], categories=["c", "b", "a"]),
+        }
+    )
+
+    result = parser.read_csv(StringIO(data), dtype={"b": dtype})
+    tm.assert_frame_equal(result, expected)
+
+
+def test_categorical_coerces_numeric(all_parsers):
+    parser = all_parsers
+    dtype = {"b": CategoricalDtype([1, 2, 3])}
+
+    data = "b\n1\n1\n2\n3"
+    expected = DataFrame({"b": Categorical([1, 1, 2, 3])})
+
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)
+
+
+def test_categorical_coerces_datetime(all_parsers):
+    parser = all_parsers
+    dti = pd.DatetimeIndex(["2017-01-01", "2018-01-01", "2019-01-01"], freq=None)
+    dtype = {"b": CategoricalDtype(dti)}
+
+    data = "b\n2017-01-01\n2018-01-01\n2019-01-01"
+    expected = DataFrame({"b": Categorical(dtype["b"].categories)})
+
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)
+
+
+def test_categorical_coerces_timestamp(all_parsers):
+    parser = all_parsers
+    dtype = {"b": CategoricalDtype([Timestamp("2014")])}
+
+    data = "b\n2014-01-01\n2014-01-01T00:00:00"
+    expected = DataFrame({"b": Categorical([Timestamp("2014")] * 2)})
+
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)
+
+
+def test_categorical_coerces_timedelta(all_parsers):
+    parser = all_parsers
+    dtype = {"b": CategoricalDtype(pd.to_timedelta(["1H", "2H", "3H"]))}
+
+    data = "b\n1H\n2H\n3H"
+    expected = DataFrame({"b": Categorical(dtype["b"].categories)})
+
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        "b\nTrue\nFalse\nNA\nFalse",
+        "b\ntrue\nfalse\nNA\nfalse",
+        "b\nTRUE\nFALSE\nNA\nFALSE",
+        "b\nTrue\nFalse\nNA\nFALSE",
+    ],
+)
+def test_categorical_dtype_coerces_boolean(all_parsers, data):
+    # see gh-20498
+    parser = all_parsers
+    dtype = {"b": CategoricalDtype([False, True])}
+    expected = DataFrame({"b": Categorical([True, False, None, False])})
+
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)
+
+
+def test_categorical_unexpected_categories(all_parsers):
+    parser = all_parsers
+    dtype = {"b": CategoricalDtype(["a", "b", "d", "e"])}
+
+    data = "b\nd\na\nc\nd"  # Unexpected c
+    expected = DataFrame({"b": Categorical(list("dacd"), dtype=dtype["b"])})
+
+    result = parser.read_csv(StringIO(data), dtype=dtype)
+    tm.assert_frame_equal(result, expected)