1
1
import numpy as np
2
2
import pytest
3
3
4
+ from pandas .errors import SettingWithCopyWarning
5
+
6
+ import pandas as pd
4
7
from pandas import (
5
8
DataFrame ,
6
9
Index ,
@@ -66,11 +69,11 @@ def test_copy_shallow(using_copy_on_write):
66
69
lambda df , copy : df .set_axis (["a" , "b" , "c" ], axis = "index" , copy = copy ),
67
70
lambda df , copy : df .rename_axis (index = "test" , copy = copy ),
68
71
lambda df , copy : df .rename_axis (columns = "test" , copy = copy ),
69
- # lambda df, copy: df.astype({'b': ' int64' }, copy=copy),
72
+ lambda df , copy : df .astype ({"b" : " int64" }, copy = copy ),
70
73
# lambda df, copy: df.swaplevel(0, 0, copy=copy),
71
74
lambda df , copy : df .swapaxes (0 , 0 , copy = copy ),
72
75
lambda df , copy : df .truncate (0 , 5 , copy = copy ),
73
- # lambda df, copy: df.infer_objects(copy=copy)
76
+ lambda df , copy : df .infer_objects (copy = copy ),
74
77
lambda df , copy : df .to_timestamp (copy = copy ),
75
78
lambda df , copy : df .to_period (freq = "D" , copy = copy ),
76
79
lambda df , copy : df .tz_localize ("US/Central" , copy = copy ),
@@ -84,11 +87,11 @@ def test_copy_shallow(using_copy_on_write):
84
87
"set_axis" ,
85
88
"rename_axis0" ,
86
89
"rename_axis1" ,
87
- # "astype", # CoW not yet implemented
90
+ "astype" ,
88
91
# "swaplevel", # only series
89
92
"swapaxes" ,
90
93
"truncate" ,
91
- # "infer_objects", # CoW not yet implemented
94
+ "infer_objects" ,
92
95
"to_timestamp" ,
93
96
"to_period" ,
94
97
"tz_localize" ,
@@ -1308,3 +1311,90 @@ def test_isetitem(using_copy_on_write):
1308
1311
assert np .shares_memory (get_array (df , "c" ), get_array (df2 , "c" ))
1309
1312
else :
1310
1313
assert not np .shares_memory (get_array (df , "c" ), get_array (df2 , "c" ))
1314
+
1315
+
1316
+ @pytest .mark .parametrize ("key" , ["a" , ["a" ]])
1317
+ def test_get (using_copy_on_write , key ):
1318
+ df = DataFrame ({"a" : [1 , 2 , 3 ], "b" : [4 , 5 , 6 ]})
1319
+ df_orig = df .copy ()
1320
+
1321
+ result = df .get (key )
1322
+
1323
+ if using_copy_on_write :
1324
+ assert np .shares_memory (get_array (result , "a" ), get_array (df , "a" ))
1325
+ result .iloc [0 ] = 0
1326
+ assert not np .shares_memory (get_array (result , "a" ), get_array (df , "a" ))
1327
+ tm .assert_frame_equal (df , df_orig )
1328
+ else :
1329
+ # for non-CoW it depends on whether we got a Series or DataFrame if it
1330
+ # is a view or copy or triggers a warning or not
1331
+ warn = SettingWithCopyWarning if isinstance (key , list ) else None
1332
+ with pd .option_context ("chained_assignment" , "warn" ):
1333
+ with tm .assert_produces_warning (warn ):
1334
+ result .iloc [0 ] = 0
1335
+
1336
+ if isinstance (key , list ):
1337
+ tm .assert_frame_equal (df , df_orig )
1338
+ else :
1339
+ assert df .iloc [0 , 0 ] == 0
1340
+
1341
+
1342
+ @pytest .mark .parametrize ("axis, key" , [(0 , 0 ), (1 , "a" )])
1343
+ @pytest .mark .parametrize (
1344
+ "dtype" , ["int64" , "float64" ], ids = ["single-block" , "mixed-block" ]
1345
+ )
1346
+ def test_xs (using_copy_on_write , using_array_manager , axis , key , dtype ):
1347
+ single_block = (dtype == "int64" ) and not using_array_manager
1348
+ is_view = single_block or (using_array_manager and axis == 1 )
1349
+ df = DataFrame (
1350
+ {"a" : [1 , 2 , 3 ], "b" : [4 , 5 , 6 ], "c" : np .array ([7 , 8 , 9 ], dtype = dtype )}
1351
+ )
1352
+ df_orig = df .copy ()
1353
+
1354
+ result = df .xs (key , axis = axis )
1355
+
1356
+ if axis == 1 or single_block :
1357
+ assert np .shares_memory (get_array (df , "a" ), get_array (result ))
1358
+ elif using_copy_on_write :
1359
+ assert result ._mgr ._has_no_reference (0 )
1360
+
1361
+ if using_copy_on_write or is_view :
1362
+ result .iloc [0 ] = 0
1363
+ else :
1364
+ with pd .option_context ("chained_assignment" , "warn" ):
1365
+ with tm .assert_produces_warning (SettingWithCopyWarning ):
1366
+ result .iloc [0 ] = 0
1367
+
1368
+ if using_copy_on_write or (not single_block and axis == 0 ):
1369
+ tm .assert_frame_equal (df , df_orig )
1370
+ else :
1371
+ assert df .iloc [0 , 0 ] == 0
1372
+
1373
+
1374
+ @pytest .mark .parametrize ("axis" , [0 , 1 ])
1375
+ @pytest .mark .parametrize ("key, level" , [("l1" , 0 ), (2 , 1 )])
1376
+ def test_xs_multiindex (using_copy_on_write , using_array_manager , key , level , axis ):
1377
+ arr = np .arange (18 ).reshape (6 , 3 )
1378
+ index = MultiIndex .from_product ([["l1" , "l2" ], [1 , 2 , 3 ]], names = ["lev1" , "lev2" ])
1379
+ df = DataFrame (arr , index = index , columns = list ("abc" ))
1380
+ if axis == 1 :
1381
+ df = df .transpose ().copy ()
1382
+ df_orig = df .copy ()
1383
+
1384
+ result = df .xs (key , level = level , axis = axis )
1385
+
1386
+ if level == 0 :
1387
+ assert np .shares_memory (
1388
+ get_array (df , df .columns [0 ]), get_array (result , result .columns [0 ])
1389
+ )
1390
+
1391
+ warn = (
1392
+ SettingWithCopyWarning
1393
+ if not using_copy_on_write and not using_array_manager
1394
+ else None
1395
+ )
1396
+ with pd .option_context ("chained_assignment" , "warn" ):
1397
+ with tm .assert_produces_warning (warn ):
1398
+ result .iloc [0 , 0 ] = 0
1399
+
1400
+ tm .assert_frame_equal (df , df_orig )
0 commit comments