@@ -1090,11 +1090,6 @@ def test_where(self):
1090
1090
expected = Series ([0 ,2 ])
1091
1091
assert_series_equal (s ,expected )
1092
1092
1093
- s = Series ([1 ,2 ])
1094
- s [[True , False ]] = [0 ]
1095
- expected = Series ([0 ,2 ])
1096
- assert_series_equal (s ,expected )
1097
-
1098
1093
# failures
1099
1094
self .assertRaises (ValueError , s .__setitem__ , tuple ([[[True , False ]]]), [0 ,2 ,3 ])
1100
1095
self .assertRaises (ValueError , s .__setitem__ , tuple ([[[True , False ]]]), [])
@@ -1142,6 +1137,24 @@ def test_where(self):
1142
1137
s = Series (np .arange (10 ))
1143
1138
mask = s > 5
1144
1139
self .assertRaises (ValueError , s .__setitem__ , mask , ([0 ]* 5 ,))
1140
+
1141
+ def test_where_broadcast (self ):
1142
+ # Test a variety of differently sized series
1143
+ for size in range (2 , 6 ):
1144
+ # Test a variety of boolean indices
1145
+ for selection in [np .resize ([True , False , False , False , False ], size ), # First element should be set
1146
+ np .resize ([True , False ], size ), # Set alternating elements]
1147
+ np .resize ([False ], size )]: # No element should be set
1148
+ # Test a variety of different numbers as content
1149
+ for item in [2.0 , np .nan , np .finfo (np .float ).max , np .finfo (np .float ).min ]:
1150
+ # Test numpy arrays, lists and tuples as the input to be broadcast
1151
+ for arr in [np .array ([item ]), [item ], (item ,)]:
1152
+ data = np .arange (size , dtype = float )
1153
+ s = Series (data )
1154
+ s [selection ] = arr
1155
+ # Construct the expected series by taking the source data or item based on the selection
1156
+ expected = Series ([item if use_item else data [i ] for i , use_item in enumerate (selection )])
1157
+ assert_series_equal (s ,expected )
1145
1158
1146
1159
def test_where_inplace (self ):
1147
1160
s = Series (np .random .randn (5 ))
0 commit comments