@@ -226,58 +226,54 @@ def test_take_nan_first_datetime(self):
226
226
expected = DatetimeIndex ([index [- 1 ], index [0 ], index [1 ]])
227
227
tm .assert_index_equal (result , expected )
228
228
229
- def test_take (self ):
229
+ @pytest .mark .parametrize ("tz" , [None , "Asia/Tokyo" ])
230
+ def test_take (self , tz ):
230
231
# GH#10295
231
- idx1 = date_range ("2011-01-01" , "2011-01-31" , freq = "D" , name = "idx" )
232
- idx2 = date_range (
233
- "2011-01-01" , "2011-01-31" , freq = "D" , tz = "Asia/Tokyo" , name = "idx"
232
+ idx = date_range ("2011-01-01" , "2011-01-31" , freq = "D" , name = "idx" , tz = tz )
233
+
234
+ result = idx .take ([0 ])
235
+ assert result == Timestamp ("2011-01-01" , tz = idx .tz )
236
+
237
+ result = idx .take ([0 , 1 , 2 ])
238
+ expected = date_range (
239
+ "2011-01-01" , "2011-01-03" , freq = "D" , tz = idx .tz , name = "idx"
234
240
)
241
+ tm .assert_index_equal (result , expected )
242
+ assert result .freq == expected .freq
235
243
236
- for idx in [idx1 , idx2 ]:
237
- result = idx .take ([0 ])
238
- assert result == Timestamp ("2011-01-01" , tz = idx .tz )
244
+ result = idx .take ([0 , 2 , 4 ])
245
+ expected = date_range (
246
+ "2011-01-01" , "2011-01-05" , freq = "2D" , tz = idx .tz , name = "idx"
247
+ )
248
+ tm .assert_index_equal (result , expected )
249
+ assert result .freq == expected .freq
239
250
240
- result = idx .take ([0 , 1 , 2 ])
241
- expected = date_range (
242
- "2011-01-01 " , "2011-01-03 " , freq = "D " , tz = idx .tz , name = "idx"
243
- )
244
- tm .assert_index_equal (result , expected )
245
- assert result .freq == expected .freq
251
+ result = idx .take ([7 , 4 , 1 ])
252
+ expected = date_range (
253
+ "2011-01-08 " , "2011-01-02 " , freq = "-3D " , tz = idx .tz , name = "idx"
254
+ )
255
+ tm .assert_index_equal (result , expected )
256
+ assert result .freq == expected .freq
246
257
247
- result = idx .take ([0 , 2 , 4 ])
248
- expected = date_range (
249
- "2011-01-01" , "2011-01-05" , freq = "2D" , tz = idx .tz , name = "idx"
250
- )
251
- tm .assert_index_equal (result , expected )
252
- assert result .freq == expected .freq
258
+ result = idx .take ([3 , 2 , 5 ])
259
+ expected = DatetimeIndex (
260
+ ["2011-01-04" , "2011-01-03" , "2011-01-06" ],
261
+ dtype = idx .dtype ,
262
+ freq = None ,
263
+ name = "idx" ,
264
+ )
265
+ tm .assert_index_equal (result , expected )
266
+ assert result .freq is None
253
267
254
- result = idx .take ([7 , 4 , 1 ])
255
- expected = date_range (
256
- "2011-01-08" , "2011-01-02" , freq = "-3D" , tz = idx .tz , name = "idx"
257
- )
258
- tm .assert_index_equal (result , expected )
259
- assert result .freq == expected .freq
260
-
261
- result = idx .take ([3 , 2 , 5 ])
262
- expected = DatetimeIndex (
263
- ["2011-01-04" , "2011-01-03" , "2011-01-06" ],
264
- dtype = idx .dtype ,
265
- freq = None ,
266
- name = "idx" ,
267
- )
268
- tm .assert_index_equal (result , expected )
269
- assert result .freq is None
270
-
271
- result = idx .take ([- 3 , 2 , 5 ])
272
- expected = DatetimeIndex (
273
- ["2011-01-29" , "2011-01-03" , "2011-01-06" ],
274
- dtype = idx .dtype ,
275
- freq = None ,
276
- tz = idx .tz ,
277
- name = "idx" ,
278
- )
279
- tm .assert_index_equal (result , expected )
280
- assert result .freq is None
268
+ result = idx .take ([- 3 , 2 , 5 ])
269
+ expected = DatetimeIndex (
270
+ ["2011-01-29" , "2011-01-03" , "2011-01-06" ],
271
+ dtype = idx .dtype ,
272
+ freq = None ,
273
+ name = "idx" ,
274
+ )
275
+ tm .assert_index_equal (result , expected )
276
+ assert result .freq is None
281
277
282
278
def test_take_invalid_kwargs (self ):
283
279
idx = date_range ("2011-01-01" , "2011-01-31" , freq = "D" , name = "idx" )
0 commit comments