|
2 | 2 |
|
3 | 3 | import pytest
|
4 | 4 |
|
5 |
| -from datetime import datetime |
| 5 | +from datetime import datetime, date |
6 | 6 | import numpy as np
|
7 | 7 | import pandas as pd
|
| 8 | +import operator as op |
8 | 9 |
|
9 | 10 | from pandas import (DatetimeIndex, Series, DataFrame,
|
10 | 11 | date_range, Index, Timedelta, Timestamp)
|
@@ -330,3 +331,21 @@ def test_loc_datetime_length_one(self):
|
330 | 331 |
|
331 | 332 | result = df.loc['2016-10-01T00:00:00':]
|
332 | 333 | tm.assert_frame_equal(result, df)
|
| 334 | + |
| 335 | + @pytest.mark.parametrize('datetimelike', [ |
| 336 | + Timestamp('20130101'), datetime(2013, 1, 1), |
| 337 | + date(2013, 1, 1), np.datetime64('2013-01-01T00:00', 'ns')]) |
| 338 | + @pytest.mark.parametrize('op,expected', [ |
| 339 | + (op.lt, [True, False, False, False]), |
| 340 | + (op.le, [True, True, False, False]), |
| 341 | + (op.eq, [False, True, False, False]), |
| 342 | + (op.gt, [False, False, False, True])]) |
| 343 | + def test_selection_by_datetimelike(self, datetimelike, op, expected): |
| 344 | + # GH issue #17965, test for ability to compare datetime64[ns] columns |
| 345 | + # to datetimelike |
| 346 | + df = DataFrame({'A': [pd.Timestamp('20120101'), |
| 347 | + pd.Timestamp('20130101'), |
| 348 | + np.nan, pd.Timestamp('20130103')]}) |
| 349 | + result = op(df.A, datetimelike) |
| 350 | + expected = Series(expected, name='A') |
| 351 | + tm.assert_series_equal(result, expected) |
0 commit comments