|
1 | 1 | import pytest
|
2 |
| -import numpy as np |
3 | 2 | import pandas as pd
|
4 |
| -import pandas._testing as tm |
5 | 3 |
|
6 | 4 |
|
7 |
| -@pytest.mark.parametrize('test_dtype', [object, 'int8', 'int16', 'int32', 'int64']) |
| 5 | +@pytest.mark.parametrize('test_dtype', [object, 'int64']) |
8 | 6 | def test_dtypes(test_dtype):
|
9 | 7 | df = pd.DataFrame({'A': pd.Series([1, 2, 3], dtype=test_dtype), 'B': [1, 2, 3]})
|
10 | 8 | expected = df.dtypes.values[0].type
|
11 | 9 |
|
12 | 10 | result = df.set_index('A').index.dtype.type
|
13 | 11 | assert result == expected
|
| 12 | + |
| 13 | + |
| 14 | +@pytest.fixture |
| 15 | +def mixed_series(): |
| 16 | + return pd.Series([1, 2, 3, 'apple', 'corn'], dtype=object) |
| 17 | + |
| 18 | + |
| 19 | +@pytest.fixture |
| 20 | +def int_series(): |
| 21 | + return pd.Series([100, 200, 300, 400, 500]) |
| 22 | + |
| 23 | + |
| 24 | +def test_dtypes_between_queries(mixed_series, int_series) |
| 25 | + df = pd.DataFrame({'item': mixed_series, 'cost': int_series}) |
| 26 | + |
| 27 | + orig_dtypes = df.dtypes |
| 28 | + item_dtype = orig_dtypes.get('item').type |
| 29 | + cost_dtype = orig_dtypes.get('cost').type |
| 30 | + expected = {'item': item_dtype, 'cost': cost_dtype} |
| 31 | + |
| 32 | + # after applying a query that would remove strings from the 'item' series with dtype: object, |
| 33 | + # that series should remain as dtype: object as it becomes an index, and again as it becomes |
| 34 | + # a column again after calling reset_index() |
| 35 | + dtypes_transformed = df.query('cost < 400').set_index('item').reset_index().dtypes |
| 36 | + item_dtype_transformed = dtypes_transformed.get('item').type |
| 37 | + cost_dtype_transformed = dtypes_transformed.get('cost').type |
| 38 | + result = {'item': item_dtype_transformed, 'cost': cost_dtype_transformed} |
| 39 | + |
| 40 | + assert result == expected |
0 commit comments