|
1 | 1 | import numpy as np |
2 | 2 | import pytest |
3 | 3 |
|
| 4 | +from pandas.errors import Pandas4Warning |
| 5 | + |
4 | 6 | from pandas.core.dtypes.dtypes import ExtensionDtype |
5 | 7 |
|
6 | 8 | import pandas as pd |
@@ -102,7 +104,12 @@ def test_select_dtypes_include_using_list_like(self, using_infer_string): |
102 | 104 | ri = df.select_dtypes(include=[str]) |
103 | 105 | tm.assert_frame_equal(ri, ei) |
104 | 106 |
|
105 | | - ri = df.select_dtypes(include=["object"]) |
| 107 | + msg = "For backward compatibility, 'str' dtypes are included" |
| 108 | + warn = None |
| 109 | + if using_infer_string: |
| 110 | + warn = Pandas4Warning |
| 111 | + with tm.assert_produces_warning(warn, match=msg): |
| 112 | + ri = df.select_dtypes(include=["object"]) |
106 | 113 | ei = df[["a"]] |
107 | 114 | tm.assert_frame_equal(ri, ei) |
108 | 115 |
|
@@ -312,15 +319,18 @@ def test_select_dtypes_not_an_attr_but_still_valid_dtype(self, using_infer_strin |
312 | 319 | ) |
313 | 320 | df["g"] = df.f.diff() |
314 | 321 | assert not hasattr(np, "u8") |
315 | | - r = df.select_dtypes(include=["i8", "O"], exclude=["timedelta"]) |
316 | | - # if using_infer_string: |
317 | | - # TODO warn |
| 322 | + |
| 323 | + msg = "For backward compatibility, 'str' dtypes are included" |
| 324 | + warn = None |
| 325 | + if using_infer_string: |
| 326 | + warn = Pandas4Warning |
| 327 | + with tm.assert_produces_warning(warn, match=msg): |
| 328 | + r = df.select_dtypes(include=["i8", "O"], exclude=["timedelta"]) |
318 | 329 | e = df[["a", "b"]] |
319 | 330 | tm.assert_frame_equal(r, e) |
320 | 331 |
|
321 | | - r = df.select_dtypes(include=["i8", "O", "timedelta64[ns]"]) |
322 | | - # if using_infer_string: |
323 | | - # TODO warn |
| 332 | + with tm.assert_produces_warning(warn, match=msg): |
| 333 | + r = df.select_dtypes(include=["i8", "O", "timedelta64[ns]"]) |
324 | 334 | e = df[["a", "b", "g"]] |
325 | 335 | tm.assert_frame_equal(r, e) |
326 | 336 |
|
@@ -497,7 +507,12 @@ def test_select_dtype_object_and_str(self, using_infer_string): |
497 | 507 | ) |
498 | 508 |
|
499 | 509 | # with "object" -> only select the object or default str dtype column |
500 | | - result = df.select_dtypes(include=["object"]) |
| 510 | + msg = "For backward compatibility, 'str' dtypes are included" |
| 511 | + warn = None |
| 512 | + if using_infer_string: |
| 513 | + warn = Pandas4Warning |
| 514 | + with tm.assert_produces_warning(warn, match=msg): |
| 515 | + result = df.select_dtypes(include=["object"]) |
501 | 516 | expected = df[["a"]] |
502 | 517 | tm.assert_frame_equal(result, expected) |
503 | 518 |
|
|
0 commit comments