Skip to content

Commit 07801be

Browse files
howsiweiSi Wei How
authored and
Si Wei How
committed
Update code and add tests
1 parent 7bc02ae commit 07801be

File tree

3 files changed

+47
-8
lines changed

3 files changed

+47
-8
lines changed

pandas/core/indexing.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -472,7 +472,8 @@ def _setitem_with_indexer(self, indexer, value):
472472
value = self._align_series(indexer, value)
473473

474474
if has_missing_columns:
475-
labels = [idx if isinstance(idx, dict) else item_labels[idx]
475+
labels = [idx['key'] if isinstance(idx, dict) else
476+
item_labels[idx]
476477
for idx in indexer[info_axis]]
477478
else:
478479
info_idx = indexer[info_axis]
@@ -527,7 +528,7 @@ def _setitem_with_indexer(self, indexer, value):
527528
lplane_indexer = 0
528529

529530
def setter(item, v):
530-
item, missing = convert_missing_indexer(item)
531+
missing = item not in self.obj.columns
531532
if not missing:
532533
s = self.obj[item]
533534
pi = plane_indexer[0] if lplane_indexer == 1 else plane_indexer

pandas/tests/frame/test_indexing.py

+9-6
Original file line numberDiff line numberDiff line change
@@ -189,33 +189,36 @@ def test_setitem_list_of_tuples(self):
189189
def test_setitem_list_all_missing_columns_scalar(self):
190190
df = self.frame.copy()
191191
df[['E', 'F']] = 1
192-
193192
df2 = self.frame.copy()
194193
df2['E'] = 1
195194
df2['F'] = 1
196-
197195
assert_frame_equal(df, df2)
198196

199197
def test_setitem_list_some_missing_columns_list(self):
200198
df = self.frame.copy()
201199
df[['A', 'E']] = [1, 2]
202-
203200
df2 = self.frame.copy()
204201
df2['A'] = 1
205202
df2['E'] = 2
206-
207203
assert_frame_equal(df, df2)
208204

209205
def test_setitem_list_some_missing_columns_dataframe(self):
210206
df = self.frame.copy()
211207
df[['A', 'E']] = self.frame[['B', 'C']]
212-
213208
df2 = self.frame.copy()
214209
df2['A'] = self.frame['B']
215210
df2['E'] = self.frame['C']
216-
217211
assert_frame_equal(df, df2)
218212

213+
def test_setitem_list_some_missing_columns_2dlist(self):
214+
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
215+
df[['B', 'C', 'D']] = [[7, 8, 9], [10, 11, 12], [13, 14, 15]]
216+
expected = pd.DataFrame([[1, 7, 8, 9],
217+
[3, 10, 11, 12],
218+
[5, 13, 14, 15]],
219+
columns=['A', 'B', 'C', 'D'])
220+
tm.assert_frame_equal(df, expected)
221+
219222
def test_setitem_mulit_index(self):
220223
# GH7655, test that assigning to a sub-frame of a frame
221224
# with multi-index columns aligns both rows and columns

pandas/tests/indexing/test_loc.py

+35
Original file line numberDiff line numberDiff line change
@@ -538,6 +538,41 @@ def test_loc_setitem_with_scalar_index(self, indexer, value):
538538

539539
assert is_scalar(result) and result == 'Z'
540540

541+
def test_loc_setitem_missing_columns_scalar_index_list_value(self):
542+
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
543+
df.loc[1, ['C', 'D']] = [7, 8]
544+
expected = pd.DataFrame([[1, 2, np.nan, np.nan],
545+
[3, 4, 7, 8],
546+
[5, 6, np.nan, np.nan]],
547+
columns=['A', 'B', 'C', 'D'])
548+
tm.assert_frame_equal(df, expected)
549+
550+
def test_loc_setitem_missing_columns_full_index_dataframe_value(self):
551+
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
552+
df2 = pd.DataFrame([[7, 8], [9, 10], [11, 12]], columns=['A', 'C'])
553+
df.loc[:, ['A', 'C']] = df2
554+
expected = pd.DataFrame([[7, 2, 8], [9, 4, 10], [11, 6, 12]],
555+
columns=['A', 'B', 'C'])
556+
tm.assert_frame_equal(df, expected)
557+
558+
def test_loc_setitem_missing_columns_list_index_scalar_value(self):
559+
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
560+
df.loc[[0, 2], ['B', 'C', 'D']] = 9
561+
expected = pd.DataFrame([[1, 9, 9, 9],
562+
[3, 4, np.nan, np.nan],
563+
[5, 9, 9, 9]],
564+
columns=['A', 'B', 'C', 'D'])
565+
tm.assert_frame_equal(df, expected)
566+
567+
def test_loc_setitem_missing_columns_range_index_2dlist_value(self):
568+
df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
569+
df.loc[1:3, ['B', 'C', 'D']] = [[7, 8, 9], [10, 11, 12]]
570+
expected = pd.DataFrame([[1, 2, np.nan, np.nan],
571+
[3, 7, 8, 9],
572+
[5, 10, 11, 12]],
573+
columns=['A', 'B', 'C', 'D'])
574+
tm.assert_frame_equal(df, expected)
575+
541576
def test_loc_coerceion(self):
542577

543578
# 12411

0 commit comments

Comments
 (0)