File tree 2 files changed +30
-2
lines changed
2 files changed +30
-2
lines changed Original file line number Diff line number Diff line change @@ -369,8 +369,16 @@ def where(
369
369
res_values = np .where (mask , self ._ndarray , value )
370
370
return self ._from_backing_data (res_values )
371
371
372
- def delete (self : NDArrayBackedExtensionArrayT , loc ) -> NDArrayBackedExtensionArrayT :
373
- res_values = np .delete (self ._ndarray , loc )
372
+ def delete (
373
+ self : NDArrayBackedExtensionArrayT , loc , axis : int = 0
374
+ ) -> NDArrayBackedExtensionArrayT :
375
+ res_values = np .delete (self ._ndarray , loc , axis = axis )
376
+ return self ._from_backing_data (res_values )
377
+
378
+ def swapaxes (
379
+ self : NDArrayBackedExtensionArrayT , axis1 , axis2
380
+ ) -> NDArrayBackedExtensionArrayT :
381
+ res_values = self ._ndarray .swapaxes (axis1 , axis2 )
374
382
return self ._from_backing_data (res_values )
375
383
376
384
# ------------------------------------------------------------------------
Original file line number Diff line number Diff line change @@ -25,6 +25,26 @@ def maybe_xfail_masked_reductions(arr, request):
25
25
26
26
27
27
class Dim2CompatTests (BaseExtensionTests ):
28
+ def test_swapaxes (self , data ):
29
+ arr2d = data .repeat (2 ).reshape (- 1 , 2 )
30
+
31
+ result = arr2d .swapaxes (0 , 1 )
32
+ expected = arr2d .T
33
+ self .assert_extension_array_equal (result , expected )
34
+
35
+ def test_delete_2d (self , data ):
36
+ arr2d = data .repeat (3 ).reshape (- 1 , 3 )
37
+
38
+ # axis = 0
39
+ result = arr2d .delete (1 , axis = 0 )
40
+ expected = data .delete (1 ).repeat (3 ).reshape (- 1 , 3 )
41
+ self .assert_extension_array_equal (result , expected )
42
+
43
+ # axis = 1
44
+ result = arr2d .delete (1 , axis = 1 )
45
+ expected = data .repeat (2 ).reshape (- 1 , 2 )
46
+ self .assert_extension_array_equal (result , expected )
47
+
28
48
def test_take_2d (self , data ):
29
49
arr2d = data .reshape (- 1 , 1 )
30
50
You can’t perform that action at this time.
0 commit comments