|
7 | 7 | import pandas as pd
|
8 | 8 | from pandas.core.arrays import ExtensionArray
|
9 | 9 | import pandas.util.testing as tm
|
| 10 | +from pandas.compat import StringIO |
| 11 | +from pandas.core.arrays.integer import ( |
| 12 | + Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype, UInt8Dtype, UInt16Dtype, |
| 13 | + UInt32Dtype, UInt64Dtype, integer_array, |
| 14 | +) |
| 15 | + |
| 16 | + |
| 17 | +def make_data(): |
| 18 | + return (list(range(1, 9)) + [np.nan] + list(range(10, 98)) |
| 19 | + + [np.nan] + [99, 100]) |
| 20 | + |
| 21 | + |
| 22 | +@pytest.fixture(params=[Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype, |
| 23 | + UInt8Dtype, UInt16Dtype, UInt32Dtype, UInt64Dtype]) |
| 24 | +def dtype(request): |
| 25 | + return request.param() |
| 26 | + |
| 27 | + |
| 28 | +@pytest.fixture |
| 29 | +def data(dtype): |
| 30 | + return integer_array(make_data(), dtype=dtype) |
10 | 31 |
|
11 | 32 |
|
12 | 33 | class DummyDtype(dtypes.ExtensionDtype):
|
@@ -92,3 +113,22 @@ def test_is_not_extension_array_dtype(dtype):
|
92 | 113 | def test_is_extension_array_dtype(dtype):
|
93 | 114 | assert isinstance(dtype, dtypes.ExtensionDtype)
|
94 | 115 | assert is_extension_array_dtype(dtype)
|
| 116 | + |
| 117 | + |
| 118 | +@pytest.mark.parametrize('engine', ['c', 'python']) |
| 119 | +def test_EA_types(engine): |
| 120 | + df = pd.DataFrame({'Int': pd.Series([1, 2, 3], dtype='Int64'), |
| 121 | + 'A': [1, 2, 1]}) |
| 122 | + data = df.to_csv(index=False) |
| 123 | + result = pd.read_csv(StringIO(data), dtype={'Int': Int64Dtype}, |
| 124 | + engine=engine) |
| 125 | + assert result is not None |
| 126 | + tm.assert_frame_equal(result, df) |
| 127 | + |
| 128 | + df = pd.DataFrame({'Int': pd.Series([1, 2, 3], dtype='Int8'), |
| 129 | + 'A': [1, 2, 1]}) |
| 130 | + data = df.to_csv(index=False) |
| 131 | + result = pd.read_csv(StringIO(data), dtype={'Int': 'Int8'}, |
| 132 | + engine=engine) |
| 133 | + assert result is not None |
| 134 | + tm.assert_frame_equal(result, df) |
0 commit comments