Skip to content

Commit ff441ec

Browse files
authored
ENH: copy keyword to set_axis (#47932)
* ENH: copy keyword to set_axis * GH ref * fix test * fix test * troubleshoot mypy * mypy fixup * test inplace=True/copy=False case * mypy fixup
1 parent 5d115bb commit ff441ec

File tree

5 files changed

+163
-20
lines changed

5 files changed

+163
-20
lines changed

doc/source/whatsnew/v1.5.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,7 @@ Other enhancements
295295
- :meth:`RangeIndex.union` now can return a :class:`RangeIndex` instead of a :class:`Int64Index` if the resulting values are equally spaced (:issue:`47557`, :issue:`43885`)
296296
- :meth:`DataFrame.compare` now accepts an argument ``result_names`` to allow the user to specify the result's names of both left and right DataFrame which are being compared. This is by default ``'self'`` and ``'other'`` (:issue:`44354`)
297297
- :class:`Interval` now supports checking whether one interval is contained by another interval (:issue:`46613`)
298+
- Added ``copy`` keyword to :meth:`Series.set_axis` and :meth:`DataFrame.set_axis` to allow user to set axis on a new object without necessarily copying the underlying data (:issue:`47932`)
298299
- :meth:`Series.add_suffix`, :meth:`DataFrame.add_suffix`, :meth:`Series.add_prefix` and :meth:`DataFrame.add_prefix` support a ``copy`` argument. If ``False``, the underlying data is not copied in the returned object (:issue:`47934`)
299300
- :meth:`DataFrame.set_index` now supports a ``copy`` keyword. If ``False``, the underlying data is not copied when a new :class:`DataFrame` is returned (:issue:`48043`)
300301

pandas/core/frame.py

+29-7
Original file line numberDiff line numberDiff line change
@@ -5036,17 +5036,34 @@ def align(
50365036

50375037
@overload
50385038
def set_axis(
5039-
self, labels, *, axis: Axis = ..., inplace: Literal[False] = ...
5039+
self,
5040+
labels,
5041+
*,
5042+
axis: Axis = ...,
5043+
inplace: Literal[False] = ...,
5044+
copy: bool | lib.NoDefault = ...,
50405045
) -> DataFrame:
50415046
...
50425047

50435048
@overload
5044-
def set_axis(self, labels, *, axis: Axis = ..., inplace: Literal[True]) -> None:
5049+
def set_axis(
5050+
self,
5051+
labels,
5052+
*,
5053+
axis: Axis = ...,
5054+
inplace: Literal[True],
5055+
copy: bool | lib.NoDefault = ...,
5056+
) -> None:
50455057
...
50465058

50475059
@overload
50485060
def set_axis(
5049-
self, labels, *, axis: Axis = ..., inplace: bool = ...
5061+
self,
5062+
labels,
5063+
*,
5064+
axis: Axis = ...,
5065+
inplace: bool = ...,
5066+
copy: bool | lib.NoDefault = ...,
50505067
) -> DataFrame | None:
50515068
...
50525069

@@ -5091,10 +5108,15 @@ def set_axis(
50915108
see_also_sub=" or columns",
50925109
)
50935110
@Appender(NDFrame.set_axis.__doc__)
5094-
def set_axis( # type: ignore[override]
5095-
self, labels, axis: Axis = 0, inplace: bool = False
5096-
) -> DataFrame | None:
5097-
return super().set_axis(labels, axis=axis, inplace=inplace)
5111+
def set_axis(
5112+
self,
5113+
labels,
5114+
axis: Axis = 0,
5115+
inplace: bool = False,
5116+
*,
5117+
copy: bool | lib.NoDefault = lib.no_default,
5118+
):
5119+
return super().set_axis(labels, axis=axis, inplace=inplace, copy=copy)
50985120

50995121
@Substitution(**_shared_doc_kwargs)
51005122
@Appender(NDFrame.reindex.__doc__)

pandas/core/generic.py

+44-8
Original file line numberDiff line numberDiff line change
@@ -711,23 +711,45 @@ def size(self) -> int:
711711

712712
@overload
713713
def set_axis(
714-
self: NDFrameT, labels, *, axis: Axis = ..., inplace: Literal[False] = ...
714+
self: NDFrameT,
715+
labels,
716+
*,
717+
axis: Axis = ...,
718+
inplace: Literal[False] = ...,
719+
copy: bool_t | lib.NoDefault = ...,
715720
) -> NDFrameT:
716721
...
717722

718723
@overload
719-
def set_axis(self, labels, *, axis: Axis = ..., inplace: Literal[True]) -> None:
724+
def set_axis(
725+
self,
726+
labels,
727+
*,
728+
axis: Axis = ...,
729+
inplace: Literal[True],
730+
copy: bool_t | lib.NoDefault = ...,
731+
) -> None:
720732
...
721733

722734
@overload
723735
def set_axis(
724-
self: NDFrameT, labels, *, axis: Axis = ..., inplace: bool_t = ...
736+
self: NDFrameT,
737+
labels,
738+
*,
739+
axis: Axis = ...,
740+
inplace: bool_t = ...,
741+
copy: bool_t | lib.NoDefault = ...,
725742
) -> NDFrameT | None:
726743
...
727744

728745
@deprecate_nonkeyword_arguments(version=None, allowed_args=["self", "labels"])
729746
def set_axis(
730-
self: NDFrameT, labels, axis: Axis = 0, inplace: bool_t = False
747+
self: NDFrameT,
748+
labels,
749+
axis: Axis = 0,
750+
inplace: bool_t = False,
751+
*,
752+
copy: bool_t | lib.NoDefault = lib.no_default,
731753
) -> NDFrameT | None:
732754
"""
733755
Assign desired index to given axis.
@@ -747,6 +769,11 @@ def set_axis(
747769
inplace : bool, default False
748770
Whether to return a new %(klass)s instance.
749771
772+
copy : bool, default True
773+
Whether to make a copy of the underlying data.
774+
775+
.. versionadded:: 1.5.0
776+
750777
Returns
751778
-------
752779
renamed : %(klass)s or None
@@ -756,16 +783,25 @@ def set_axis(
756783
--------
757784
%(klass)s.rename_axis : Alter the name of the index%(see_also_sub)s.
758785
"""
786+
if inplace:
787+
if copy is True:
788+
raise ValueError("Cannot specify both inplace=True and copy=True")
789+
copy = False
790+
elif copy is lib.no_default:
791+
copy = True
792+
759793
self._check_inplace_and_allows_duplicate_labels(inplace)
760-
return self._set_axis_nocheck(labels, axis, inplace)
794+
return self._set_axis_nocheck(labels, axis, inplace, copy=copy)
761795

762796
@final
763-
def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool_t):
797+
def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool_t, copy: bool_t):
764798
# NDFrame.rename with inplace=False calls set_axis(inplace=True) on a copy.
765799
if inplace:
766800
setattr(self, self._get_axis_name(axis), labels)
767801
else:
768-
obj = self.copy()
802+
# With copy=False, we create a new object but don't copy the
803+
# underlying data.
804+
obj = self.copy(deep=copy)
769805
obj.set_axis(labels, axis=axis, inplace=True)
770806
return obj
771807

@@ -1053,7 +1089,7 @@ def _rename(
10531089
raise KeyError(f"{missing_labels} not found in axis")
10541090

10551091
new_index = ax._transform_index(f, level=level)
1056-
result._set_axis_nocheck(new_index, axis=axis_no, inplace=True)
1092+
result._set_axis_nocheck(new_index, axis=axis_no, inplace=True, copy=False)
10571093
result._clear_item_cache()
10581094

10591095
if inplace:

pandas/core/series.py

+26-5
Original file line numberDiff line numberDiff line change
@@ -4976,17 +4976,34 @@ def rename(
49764976

49774977
@overload
49784978
def set_axis(
4979-
self, labels, *, axis: Axis = ..., inplace: Literal[False] = ...
4979+
self,
4980+
labels,
4981+
*,
4982+
axis: Axis = ...,
4983+
inplace: Literal[False] = ...,
4984+
copy: bool | lib.NoDefault = ...,
49804985
) -> Series:
49814986
...
49824987

49834988
@overload
4984-
def set_axis(self, labels, *, axis: Axis = ..., inplace: Literal[True]) -> None:
4989+
def set_axis(
4990+
self,
4991+
labels,
4992+
*,
4993+
axis: Axis = ...,
4994+
inplace: Literal[True],
4995+
copy: bool | lib.NoDefault = ...,
4996+
) -> None:
49854997
...
49864998

49874999
@overload
49885000
def set_axis(
4989-
self, labels, *, axis: Axis = ..., inplace: bool = ...
5001+
self,
5002+
labels,
5003+
*,
5004+
axis: Axis = ...,
5005+
inplace: bool = ...,
5006+
copy: bool | lib.NoDefault = ...,
49905007
) -> Series | None:
49915008
...
49925009

@@ -5018,9 +5035,13 @@ def set_axis(
50185035
)
50195036
@Appender(NDFrame.set_axis.__doc__)
50205037
def set_axis( # type: ignore[override]
5021-
self, labels, axis: Axis = 0, inplace: bool = False
5038+
self,
5039+
labels,
5040+
axis: Axis = 0,
5041+
inplace: bool = False,
5042+
copy: bool | lib.NoDefault = lib.no_default,
50225043
) -> Series | None:
5023-
return super().set_axis(labels, axis=axis, inplace=inplace)
5044+
return super().set_axis(labels, axis=axis, inplace=inplace, copy=copy)
50245045

50255046
# error: Cannot determine type of 'reindex'
50265047
@doc(

pandas/tests/frame/methods/test_set_axis.py

+63
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,69 @@ def test_set_axis(self, obj):
2424
result = obj.set_axis(new_index, axis=0, inplace=False)
2525
tm.assert_equal(expected, result)
2626

27+
def test_set_axis_copy(self, obj):
28+
# Test copy keyword GH#47932
29+
new_index = list("abcd")[: len(obj)]
30+
31+
orig = obj.iloc[:]
32+
expected = obj.copy()
33+
expected.index = new_index
34+
35+
with pytest.raises(
36+
ValueError, match="Cannot specify both inplace=True and copy=True"
37+
):
38+
obj.set_axis(new_index, axis=0, inplace=True, copy=True)
39+
40+
result = obj.set_axis(new_index, axis=0, copy=True)
41+
tm.assert_equal(expected, result)
42+
assert result is not obj
43+
# check we DID make a copy
44+
if obj.ndim == 1:
45+
assert not tm.shares_memory(result, obj)
46+
else:
47+
assert not any(
48+
tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
49+
for i in range(obj.shape[1])
50+
)
51+
52+
result = obj.set_axis(new_index, axis=0, copy=False)
53+
tm.assert_equal(expected, result)
54+
assert result is not obj
55+
# check we did NOT make a copy
56+
if obj.ndim == 1:
57+
assert tm.shares_memory(result, obj)
58+
else:
59+
assert all(
60+
tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
61+
for i in range(obj.shape[1])
62+
)
63+
64+
# copy defaults to True
65+
result = obj.set_axis(new_index, axis=0)
66+
tm.assert_equal(expected, result)
67+
assert result is not obj
68+
# check we DID make a copy
69+
if obj.ndim == 1:
70+
assert not tm.shares_memory(result, obj)
71+
else:
72+
assert not any(
73+
tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
74+
for i in range(obj.shape[1])
75+
)
76+
77+
# Do this last since it alters obj inplace
78+
res = obj.set_axis(new_index, inplace=True, copy=False)
79+
assert res is None
80+
tm.assert_equal(expected, obj)
81+
# check we did NOT make a copy
82+
if obj.ndim == 1:
83+
assert tm.shares_memory(obj, orig)
84+
else:
85+
assert all(
86+
tm.shares_memory(obj.iloc[:, i], orig.iloc[:, i])
87+
for i in range(obj.shape[1])
88+
)
89+
2790
@pytest.mark.parametrize("axis", [0, "index", 1, "columns"])
2891
def test_set_axis_inplace_axis(self, axis, obj):
2992
# GH#14636

0 commit comments

Comments
 (0)