Skip to content

Commit 4e929c5

Browse files
howsiweihowsiwei
committed
Add docstring, types and tests
1 parent 3c5c765 commit 4e929c5

File tree

3 files changed

+113
-35
lines changed

3 files changed

+113
-35
lines changed

‎pandas/core/indexing.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -166,9 +166,18 @@ def _get_loc(self, key: int, axis: int):
166166
def _slice(self, obj, axis: int, kind=None):
167167
return self.obj._slice(obj, axis=axis, kind=kind)
168168

169-
def _ensure_listlike_indexer(self, key, axis):
170-
# if `key` is an index to multiple columns by name, add each column
171-
# if it does not already exist
169+
def _ensure_listlike_indexer(self, key, axis: int):
170+
"""
171+
Ensure that a list-like of keys are all present by adding them if they
172+
do not already exist.
173+
174+
Parameters
175+
----------
176+
key : list-like
177+
Target labels
178+
axis : int
179+
Dimension on which the indexing is being made
180+
"""
172181
if not isinstance(self.obj._get_axis(axis), ABCMultiIndex) and all(
173182
is_hashable(k) for k in key
174183
):

‎pandas/tests/frame/indexing/test_indexing.py

Lines changed: 53 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -202,28 +202,63 @@ def test_setitem_list_of_tuples(self, float_frame):
202202
expected = Series(tuples, index=float_frame.index, name="tuples")
203203
tm.assert_series_equal(result, expected)
204204

205-
@pytest.mark.parametrize("columns", [["A", "E"], ["E", "F"]])
206205
@pytest.mark.parametrize(
207-
"box",
206+
"columns,box,expected",
208207
[
209-
lambda x: 1,
210-
lambda x: [1, 2],
211-
lambda x: np.array([1, 2]),
212-
lambda x: x[["B", "C"]],
213-
lambda x: x[["B", "A"]].values,
214-
lambda x: x[["A", "C"]].values.tolist(),
208+
(
209+
["A", "B", "C", "D"],
210+
7,
211+
pd.DataFrame(
212+
[[7, 7, 7, 7], [7, 7, 7, 7], [7, 7, 7, 7]],
213+
columns=["A", "B", "C", "D"],
214+
),
215+
),
216+
(
217+
["C", "D"],
218+
[7, 8],
219+
pd.DataFrame(
220+
[[1, 2, 7, 8], [3, 4, 7, 8], [5, 6, 7, 8]],
221+
columns=["A", "B", "C", "D"],
222+
),
223+
),
224+
(
225+
["A", "B", "C"],
226+
np.array([7, 8, 9]),
227+
pd.DataFrame(
228+
[[7, 8, 9], [7, 8, 9], [7, 8, 9]],
229+
columns=["A", "B", "C"],
230+
),
231+
),
232+
(
233+
["B", "C", "D"],
234+
[[7, 8, 9], [10, 11, 12], [13, 14, 15]],
235+
pd.DataFrame(
236+
[[1, 7, 8, 9], [3, 10, 11, 12], [5, 13, 14, 15]],
237+
columns=["A", "B", "C", "D"],
238+
),
239+
),
240+
(
241+
["C", "A", "D"],
242+
np.array([[7, 8, 9], [10, 11, 12], [13, 14, 15]]),
243+
pd.DataFrame(
244+
[[8, 2, 7, 9], [11, 4, 10, 12], [14, 6, 13, 15]],
245+
columns=["A", "B", "C", "D"],
246+
),
247+
),
248+
(
249+
["A", "C"],
250+
pd.DataFrame([[7, 8], [9, 10], [11, 12]], columns=["A", "C"]),
251+
pd.DataFrame(
252+
[[7, 2, 8], [9, 4, 10], [11, 6, 12]], columns=["A", "B", "C"]
253+
),
254+
),
215255
],
216256
)
217-
def test_setitem_list_missing_columns(self, float_frame, columns, box):
218-
# GH 26534
219-
result = float_frame.copy()
220-
result[columns] = box(float_frame)
221-
expected = float_frame.copy()
222-
for col in columns:
223-
if col not in expected.columns:
224-
expected[col] = np.nan
225-
expected[columns] = box(float_frame)
226-
tm.assert_frame_equal(result, expected)
257+
def test_setitem_list_missing_columns(self, columns, box, expected):
258+
# GH 29334
259+
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
260+
df[columns] = box
261+
tm.assert_frame_equal(df, expected)
227262

228263
def test_setitem_mulit_index(self):
229264
# GH7655, test that assigning to a sub-frame of a frame

‎pandas/tests/indexing/test_loc.py

Lines changed: 48 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -751,28 +751,62 @@ def test_loc_setitem_with_scalar_index(self, indexer, value):
751751
assert is_scalar(result) and result == "Z"
752752

753753
@pytest.mark.parametrize(
754-
"index,box",
754+
"index,box,expected",
755755
[
756-
((1, ["C", "D"]), [7, 8]),
756+
(
757+
([0, 2], ["A", "B", "C", "D"]),
758+
7,
759+
pd.DataFrame(
760+
[[7, 7, 7, 7], [3, 4, np.nan, np.nan], [7, 7, 7, 7]],
761+
columns=["A", "B", "C", "D"],
762+
),
763+
),
764+
(
765+
(1, ["C", "D"]),
766+
[7, 8],
767+
pd.DataFrame(
768+
[[1, 2, np.nan, np.nan], [3, 4, 7, 8], [5, 6, np.nan, np.nan]],
769+
columns=["A", "B", "C", "D"],
770+
),
771+
),
772+
(
773+
(1, ["A", "B", "C"]),
774+
np.array([7, 8, 9]),
775+
pd.DataFrame(
776+
[[1, 2, np.nan], [7, 8, 9], [5, 6, np.nan]],
777+
columns=["A", "B", "C"],
778+
),
779+
),
780+
(
781+
(slice(1, 3, None), ["B", "C", "D"]),
782+
[[7, 8, 9], [10, 11, 12]],
783+
pd.DataFrame(
784+
[[1, 2, np.nan, np.nan], [3, 7, 8, 9], [5, 10, 11, 12]],
785+
columns=["A", "B", "C", "D"],
786+
),
787+
),
788+
(
789+
(slice(1, 3, None), ["C", "A", "D"]),
790+
np.array([[7, 8, 9], [10, 11, 12]]),
791+
pd.DataFrame(
792+
[[1, 2, np.nan, np.nan], [8, 4, 7, 9], [11, 6, 10, 12]],
793+
columns=["A", "B", "C", "D"],
794+
),
795+
),
757796
(
758797
(slice(None, None, None), ["A", "C"]),
759798
pd.DataFrame([[7, 8], [9, 10], [11, 12]], columns=["A", "C"]),
799+
pd.DataFrame(
800+
[[7, 2, 8], [9, 4, 10], [11, 6, 12]], columns=["A", "B", "C"]
801+
),
760802
),
761-
(([0, 2], ["B", "C", "D"]), 9),
762-
((slice(1, 3, None), ["B", "C", "D"]), [[7, 8, 9], [10, 11, 12]]),
763803
],
764804
)
765-
def test_loc_setitem_missing_columns(self, index, box):
766-
# GH 26534
805+
def test_loc_setitem_missing_columns(self, index, box, expected):
806+
# GH 29334
767807
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
768-
result = df.copy()
769-
result.loc[index] = box
770-
expected = df.copy()
771-
for col in index[1]:
772-
if col not in expected.columns:
773-
expected[col] = np.nan
774-
expected.loc[index] = box
775-
tm.assert_frame_equal(result, expected)
808+
df.loc[index] = box
809+
tm.assert_frame_equal(df, expected)
776810

777811
def test_loc_coercion(self):
778812

0 commit comments

Comments
 (0)