test_indexing.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815
  1. from datetime import datetime
  2. import re
  3. import numpy as np
  4. import pytest
  5. from pandas._libs.tslibs import period as libperiod
  6. from pandas.errors import InvalidIndexError
  7. import pandas as pd
  8. from pandas import (
  9. DatetimeIndex,
  10. NaT,
  11. Period,
  12. PeriodIndex,
  13. Series,
  14. Timedelta,
  15. date_range,
  16. notna,
  17. period_range,
  18. )
  19. import pandas._testing as tm
  20. dti4 = date_range("2016-01-01", periods=4)
  21. dti = dti4[:-1]
  22. rng = pd.Index(range(3))
  23. @pytest.fixture(
  24. params=[
  25. dti,
  26. dti.tz_localize("UTC"),
  27. dti.to_period("W"),
  28. dti - dti[0],
  29. rng,
  30. pd.Index([1, 2, 3]),
  31. pd.Index([2.0, 3.0, 4.0]),
  32. pd.Index([4, 5, 6], dtype="u8"),
  33. pd.IntervalIndex.from_breaks(dti4),
  34. ]
  35. )
  36. def non_comparable_idx(request):
  37. # All have length 3
  38. return request.param
  39. class TestGetItem:
  40. def test_getitem_slice_keeps_name(self):
  41. idx = period_range("20010101", periods=10, freq="D", name="bob")
  42. assert idx.name == idx[1:].name
  43. def test_getitem(self):
  44. idx1 = period_range("2011-01-01", "2011-01-31", freq="D", name="idx")
  45. for idx in [idx1]:
  46. result = idx[0]
  47. assert result == Period("2011-01-01", freq="D")
  48. result = idx[-1]
  49. assert result == Period("2011-01-31", freq="D")
  50. result = idx[0:5]
  51. expected = period_range("2011-01-01", "2011-01-05", freq="D", name="idx")
  52. tm.assert_index_equal(result, expected)
  53. assert result.freq == expected.freq
  54. assert result.freq == "D"
  55. result = idx[0:10:2]
  56. expected = PeriodIndex(
  57. ["2011-01-01", "2011-01-03", "2011-01-05", "2011-01-07", "2011-01-09"],
  58. freq="D",
  59. name="idx",
  60. )
  61. tm.assert_index_equal(result, expected)
  62. assert result.freq == expected.freq
  63. assert result.freq == "D"
  64. result = idx[-20:-5:3]
  65. expected = PeriodIndex(
  66. ["2011-01-12", "2011-01-15", "2011-01-18", "2011-01-21", "2011-01-24"],
  67. freq="D",
  68. name="idx",
  69. )
  70. tm.assert_index_equal(result, expected)
  71. assert result.freq == expected.freq
  72. assert result.freq == "D"
  73. result = idx[4::-1]
  74. expected = PeriodIndex(
  75. ["2011-01-05", "2011-01-04", "2011-01-03", "2011-01-02", "2011-01-01"],
  76. freq="D",
  77. name="idx",
  78. )
  79. tm.assert_index_equal(result, expected)
  80. assert result.freq == expected.freq
  81. assert result.freq == "D"
  82. def test_getitem_index(self):
  83. idx = period_range("2007-01", periods=10, freq="M", name="x")
  84. result = idx[[1, 3, 5]]
  85. exp = PeriodIndex(["2007-02", "2007-04", "2007-06"], freq="M", name="x")
  86. tm.assert_index_equal(result, exp)
  87. result = idx[[True, True, False, False, False, True, True, False, False, False]]
  88. exp = PeriodIndex(
  89. ["2007-01", "2007-02", "2007-06", "2007-07"], freq="M", name="x"
  90. )
  91. tm.assert_index_equal(result, exp)
  92. def test_getitem_partial(self):
  93. rng = period_range("2007-01", periods=50, freq="M")
  94. ts = Series(np.random.randn(len(rng)), rng)
  95. with pytest.raises(KeyError, match=r"^'2006'$"):
  96. ts["2006"]
  97. result = ts["2008"]
  98. assert (result.index.year == 2008).all()
  99. result = ts["2008":"2009"]
  100. assert len(result) == 24
  101. result = ts["2008-1":"2009-12"]
  102. assert len(result) == 24
  103. result = ts["2008Q1":"2009Q4"]
  104. assert len(result) == 24
  105. result = ts[:"2009"]
  106. assert len(result) == 36
  107. result = ts["2009":]
  108. assert len(result) == 50 - 24
  109. exp = result
  110. result = ts[24:]
  111. tm.assert_series_equal(exp, result)
  112. ts = pd.concat([ts[10:], ts[10:]])
  113. msg = "left slice bound for non-unique label: '2008'"
  114. with pytest.raises(KeyError, match=msg):
  115. ts[slice("2008", "2009")]
  116. def test_getitem_datetime(self):
  117. rng = period_range(start="2012-01-01", periods=10, freq="W-MON")
  118. ts = Series(range(len(rng)), index=rng)
  119. dt1 = datetime(2011, 10, 2)
  120. dt4 = datetime(2012, 4, 20)
  121. rs = ts[dt1:dt4]
  122. tm.assert_series_equal(rs, ts)
  123. def test_getitem_nat(self):
  124. idx = PeriodIndex(["2011-01", "NaT", "2011-02"], freq="M")
  125. assert idx[0] == Period("2011-01", freq="M")
  126. assert idx[1] is NaT
  127. s = Series([0, 1, 2], index=idx)
  128. assert s[NaT] == 1
  129. s = Series(idx, index=idx)
  130. assert s[Period("2011-01", freq="M")] == Period("2011-01", freq="M")
  131. assert s[NaT] is NaT
  132. def test_getitem_list_periods(self):
  133. # GH 7710
  134. rng = period_range(start="2012-01-01", periods=10, freq="D")
  135. ts = Series(range(len(rng)), index=rng)
  136. exp = ts.iloc[[1]]
  137. tm.assert_series_equal(ts[[Period("2012-01-02", freq="D")]], exp)
  138. @pytest.mark.arm_slow
  139. def test_getitem_seconds(self):
  140. # GH#6716
  141. didx = date_range(start="2013/01/01 09:00:00", freq="S", periods=4000)
  142. pidx = period_range(start="2013/01/01 09:00:00", freq="S", periods=4000)
  143. for idx in [didx, pidx]:
  144. # getitem against index should raise ValueError
  145. values = [
  146. "2014",
  147. "2013/02",
  148. "2013/01/02",
  149. "2013/02/01 9H",
  150. "2013/02/01 09:00",
  151. ]
  152. for val in values:
  153. # GH7116
  154. # these show deprecations as we are trying
  155. # to slice with non-integer indexers
  156. with pytest.raises(IndexError, match="only integers, slices"):
  157. idx[val]
  158. ser = Series(np.random.rand(len(idx)), index=idx)
  159. tm.assert_series_equal(ser["2013/01/01 10:00"], ser[3600:3660])
  160. tm.assert_series_equal(ser["2013/01/01 9H"], ser[:3600])
  161. for d in ["2013/01/01", "2013/01", "2013"]:
  162. tm.assert_series_equal(ser[d], ser)
  163. @pytest.mark.parametrize(
  164. "idx_range",
  165. [
  166. date_range,
  167. period_range,
  168. ],
  169. )
  170. def test_getitem_day(self, idx_range):
  171. # GH#6716
  172. # Confirm DatetimeIndex and PeriodIndex works identically
  173. # getitem against index should raise ValueError
  174. idx = idx_range(start="2013/01/01", freq="D", periods=400)
  175. values = [
  176. "2014",
  177. "2013/02",
  178. "2013/01/02",
  179. "2013/02/01 9H",
  180. "2013/02/01 09:00",
  181. ]
  182. for val in values:
  183. # GH7116
  184. # these show deprecations as we are trying
  185. # to slice with non-integer indexers
  186. with pytest.raises(IndexError, match="only integers, slices"):
  187. idx[val]
  188. ser = Series(np.random.rand(len(idx)), index=idx)
  189. tm.assert_series_equal(ser["2013/01"], ser[0:31])
  190. tm.assert_series_equal(ser["2013/02"], ser[31:59])
  191. tm.assert_series_equal(ser["2014"], ser[365:])
  192. invalid = ["2013/02/01 9H", "2013/02/01 09:00"]
  193. for val in invalid:
  194. with pytest.raises(KeyError, match=val):
  195. ser[val]
  196. class TestGetLoc:
  197. def test_get_loc_msg(self):
  198. idx = period_range("2000-1-1", freq="A", periods=10)
  199. bad_period = Period("2012", "A")
  200. with pytest.raises(KeyError, match=r"^Period\('2012', 'A-DEC'\)$"):
  201. idx.get_loc(bad_period)
  202. try:
  203. idx.get_loc(bad_period)
  204. except KeyError as inst:
  205. assert inst.args[0] == bad_period
  206. def test_get_loc_nat(self):
  207. didx = DatetimeIndex(["2011-01-01", "NaT", "2011-01-03"])
  208. pidx = PeriodIndex(["2011-01-01", "NaT", "2011-01-03"], freq="M")
  209. # check DatetimeIndex compat
  210. for idx in [didx, pidx]:
  211. assert idx.get_loc(NaT) == 1
  212. assert idx.get_loc(None) == 1
  213. assert idx.get_loc(float("nan")) == 1
  214. assert idx.get_loc(np.nan) == 1
  215. def test_get_loc(self):
  216. # GH 17717
  217. p0 = Period("2017-09-01")
  218. p1 = Period("2017-09-02")
  219. p2 = Period("2017-09-03")
  220. # get the location of p1/p2 from
  221. # monotonic increasing PeriodIndex with non-duplicate
  222. idx0 = PeriodIndex([p0, p1, p2])
  223. expected_idx1_p1 = 1
  224. expected_idx1_p2 = 2
  225. assert idx0.get_loc(p1) == expected_idx1_p1
  226. assert idx0.get_loc(str(p1)) == expected_idx1_p1
  227. assert idx0.get_loc(p2) == expected_idx1_p2
  228. assert idx0.get_loc(str(p2)) == expected_idx1_p2
  229. msg = "Cannot interpret 'foo' as period"
  230. with pytest.raises(KeyError, match=msg):
  231. idx0.get_loc("foo")
  232. with pytest.raises(KeyError, match=r"^1\.1$"):
  233. idx0.get_loc(1.1)
  234. with pytest.raises(InvalidIndexError, match=re.escape(str(idx0))):
  235. idx0.get_loc(idx0)
  236. # get the location of p1/p2 from
  237. # monotonic increasing PeriodIndex with duplicate
  238. idx1 = PeriodIndex([p1, p1, p2])
  239. expected_idx1_p1 = slice(0, 2)
  240. expected_idx1_p2 = 2
  241. assert idx1.get_loc(p1) == expected_idx1_p1
  242. assert idx1.get_loc(str(p1)) == expected_idx1_p1
  243. assert idx1.get_loc(p2) == expected_idx1_p2
  244. assert idx1.get_loc(str(p2)) == expected_idx1_p2
  245. msg = "Cannot interpret 'foo' as period"
  246. with pytest.raises(KeyError, match=msg):
  247. idx1.get_loc("foo")
  248. with pytest.raises(KeyError, match=r"^1\.1$"):
  249. idx1.get_loc(1.1)
  250. with pytest.raises(InvalidIndexError, match=re.escape(str(idx1))):
  251. idx1.get_loc(idx1)
  252. # get the location of p1/p2 from
  253. # non-monotonic increasing/decreasing PeriodIndex with duplicate
  254. idx2 = PeriodIndex([p2, p1, p2])
  255. expected_idx2_p1 = 1
  256. expected_idx2_p2 = np.array([True, False, True])
  257. assert idx2.get_loc(p1) == expected_idx2_p1
  258. assert idx2.get_loc(str(p1)) == expected_idx2_p1
  259. tm.assert_numpy_array_equal(idx2.get_loc(p2), expected_idx2_p2)
  260. tm.assert_numpy_array_equal(idx2.get_loc(str(p2)), expected_idx2_p2)
  261. def test_get_loc_integer(self):
  262. dti = date_range("2016-01-01", periods=3)
  263. pi = dti.to_period("D")
  264. with pytest.raises(KeyError, match="16801"):
  265. pi.get_loc(16801)
  266. pi2 = dti.to_period("Y") # duplicates, ordinals are all 46
  267. with pytest.raises(KeyError, match="46"):
  268. pi2.get_loc(46)
  269. def test_get_loc_invalid_string_raises_keyerror(self):
  270. # GH#34240
  271. pi = period_range("2000", periods=3, name="A")
  272. with pytest.raises(KeyError, match="A"):
  273. pi.get_loc("A")
  274. ser = Series([1, 2, 3], index=pi)
  275. with pytest.raises(KeyError, match="A"):
  276. ser.loc["A"]
  277. with pytest.raises(KeyError, match="A"):
  278. ser["A"]
  279. assert "A" not in ser
  280. assert "A" not in pi
  281. def test_get_loc_mismatched_freq(self):
  282. # see also test_get_indexer_mismatched_dtype testing we get analogous
  283. # behavior for get_loc
  284. dti = date_range("2016-01-01", periods=3)
  285. pi = dti.to_period("D")
  286. pi2 = dti.to_period("W")
  287. pi3 = pi.view(pi2.dtype) # i.e. matching i8 representations
  288. with pytest.raises(KeyError, match="W-SUN"):
  289. pi.get_loc(pi2[0])
  290. with pytest.raises(KeyError, match="W-SUN"):
  291. # even though we have matching i8 values
  292. pi.get_loc(pi3[0])
  293. class TestGetIndexer:
  294. def test_get_indexer(self):
  295. # GH 17717
  296. p1 = Period("2017-09-01")
  297. p2 = Period("2017-09-04")
  298. p3 = Period("2017-09-07")
  299. tp0 = Period("2017-08-31")
  300. tp1 = Period("2017-09-02")
  301. tp2 = Period("2017-09-05")
  302. tp3 = Period("2017-09-09")
  303. idx = PeriodIndex([p1, p2, p3])
  304. tm.assert_numpy_array_equal(
  305. idx.get_indexer(idx), np.array([0, 1, 2], dtype=np.intp)
  306. )
  307. target = PeriodIndex([tp0, tp1, tp2, tp3])
  308. tm.assert_numpy_array_equal(
  309. idx.get_indexer(target, "pad"), np.array([-1, 0, 1, 2], dtype=np.intp)
  310. )
  311. tm.assert_numpy_array_equal(
  312. idx.get_indexer(target, "backfill"), np.array([0, 1, 2, -1], dtype=np.intp)
  313. )
  314. tm.assert_numpy_array_equal(
  315. idx.get_indexer(target, "nearest"), np.array([0, 0, 1, 2], dtype=np.intp)
  316. )
  317. res = idx.get_indexer(target, "nearest", tolerance=Timedelta("1 day"))
  318. tm.assert_numpy_array_equal(res, np.array([0, 0, 1, -1], dtype=np.intp))
  319. def test_get_indexer_mismatched_dtype(self):
  320. # Check that we return all -1s and do not raise or cast incorrectly
  321. dti = date_range("2016-01-01", periods=3)
  322. pi = dti.to_period("D")
  323. pi2 = dti.to_period("W")
  324. expected = np.array([-1, -1, -1], dtype=np.intp)
  325. result = pi.get_indexer(dti)
  326. tm.assert_numpy_array_equal(result, expected)
  327. # This should work in both directions
  328. result = dti.get_indexer(pi)
  329. tm.assert_numpy_array_equal(result, expected)
  330. result = pi.get_indexer(pi2)
  331. tm.assert_numpy_array_equal(result, expected)
  332. # We expect the same from get_indexer_non_unique
  333. result = pi.get_indexer_non_unique(dti)[0]
  334. tm.assert_numpy_array_equal(result, expected)
  335. result = dti.get_indexer_non_unique(pi)[0]
  336. tm.assert_numpy_array_equal(result, expected)
  337. result = pi.get_indexer_non_unique(pi2)[0]
  338. tm.assert_numpy_array_equal(result, expected)
  339. def test_get_indexer_mismatched_dtype_different_length(self, non_comparable_idx):
  340. # without method we aren't checking inequalities, so get all-missing
  341. # but do not raise
  342. dti = date_range("2016-01-01", periods=3)
  343. pi = dti.to_period("D")
  344. other = non_comparable_idx
  345. res = pi[:-1].get_indexer(other)
  346. expected = -np.ones(other.shape, dtype=np.intp)
  347. tm.assert_numpy_array_equal(res, expected)
  348. @pytest.mark.parametrize("method", ["pad", "backfill", "nearest"])
  349. def test_get_indexer_mismatched_dtype_with_method(self, non_comparable_idx, method):
  350. dti = date_range("2016-01-01", periods=3)
  351. pi = dti.to_period("D")
  352. other = non_comparable_idx
  353. msg = re.escape(f"Cannot compare dtypes {pi.dtype} and {other.dtype}")
  354. with pytest.raises(TypeError, match=msg):
  355. pi.get_indexer(other, method=method)
  356. for dtype in ["object", "category"]:
  357. other2 = other.astype(dtype)
  358. if dtype == "object" and isinstance(other, PeriodIndex):
  359. continue
  360. # Two different error message patterns depending on dtypes
  361. msg = "|".join(
  362. [
  363. re.escape(msg)
  364. for msg in (
  365. f"Cannot compare dtypes {pi.dtype} and {other.dtype}",
  366. " not supported between instances of ",
  367. )
  368. ]
  369. )
  370. with pytest.raises(TypeError, match=msg):
  371. pi.get_indexer(other2, method=method)
  372. def test_get_indexer_non_unique(self):
  373. # GH 17717
  374. p1 = Period("2017-09-02")
  375. p2 = Period("2017-09-03")
  376. p3 = Period("2017-09-04")
  377. p4 = Period("2017-09-05")
  378. idx1 = PeriodIndex([p1, p2, p1])
  379. idx2 = PeriodIndex([p2, p1, p3, p4])
  380. result = idx1.get_indexer_non_unique(idx2)
  381. expected_indexer = np.array([1, 0, 2, -1, -1], dtype=np.intp)
  382. expected_missing = np.array([2, 3], dtype=np.intp)
  383. tm.assert_numpy_array_equal(result[0], expected_indexer)
  384. tm.assert_numpy_array_equal(result[1], expected_missing)
  385. # TODO: This method came from test_period; de-dup with version above
  386. def test_get_indexer2(self):
  387. idx = period_range("2000-01-01", periods=3).asfreq("H", how="start")
  388. tm.assert_numpy_array_equal(
  389. idx.get_indexer(idx), np.array([0, 1, 2], dtype=np.intp)
  390. )
  391. target = PeriodIndex(
  392. ["1999-12-31T23", "2000-01-01T12", "2000-01-02T01"], freq="H"
  393. )
  394. tm.assert_numpy_array_equal(
  395. idx.get_indexer(target, "pad"), np.array([-1, 0, 1], dtype=np.intp)
  396. )
  397. tm.assert_numpy_array_equal(
  398. idx.get_indexer(target, "backfill"), np.array([0, 1, 2], dtype=np.intp)
  399. )
  400. tm.assert_numpy_array_equal(
  401. idx.get_indexer(target, "nearest"), np.array([0, 1, 1], dtype=np.intp)
  402. )
  403. tm.assert_numpy_array_equal(
  404. idx.get_indexer(target, "nearest", tolerance="1 hour"),
  405. np.array([0, -1, 1], dtype=np.intp),
  406. )
  407. msg = "Input has different freq=None from PeriodArray\\(freq=H\\)"
  408. with pytest.raises(ValueError, match=msg):
  409. idx.get_indexer(target, "nearest", tolerance="1 minute")
  410. tm.assert_numpy_array_equal(
  411. idx.get_indexer(target, "nearest", tolerance="1 day"),
  412. np.array([0, 1, 1], dtype=np.intp),
  413. )
  414. tol_raw = [
  415. Timedelta("1 hour"),
  416. Timedelta("1 hour"),
  417. np.timedelta64(1, "D"),
  418. ]
  419. tm.assert_numpy_array_equal(
  420. idx.get_indexer(
  421. target, "nearest", tolerance=[np.timedelta64(x) for x in tol_raw]
  422. ),
  423. np.array([0, -1, 1], dtype=np.intp),
  424. )
  425. tol_bad = [
  426. Timedelta("2 hour").to_timedelta64(),
  427. Timedelta("1 hour").to_timedelta64(),
  428. np.timedelta64(1, "M"),
  429. ]
  430. with pytest.raises(
  431. libperiod.IncompatibleFrequency, match="Input has different freq=None from"
  432. ):
  433. idx.get_indexer(target, "nearest", tolerance=tol_bad)
  434. class TestWhere:
  435. def test_where(self, listlike_box):
  436. i = period_range("20130101", periods=5, freq="D")
  437. cond = [True] * len(i)
  438. expected = i
  439. result = i.where(listlike_box(cond))
  440. tm.assert_index_equal(result, expected)
  441. cond = [False] + [True] * (len(i) - 1)
  442. expected = PeriodIndex([NaT] + i[1:].tolist(), freq="D")
  443. result = i.where(listlike_box(cond))
  444. tm.assert_index_equal(result, expected)
  445. def test_where_other(self):
  446. i = period_range("20130101", periods=5, freq="D")
  447. for arr in [np.nan, NaT]:
  448. result = i.where(notna(i), other=arr)
  449. expected = i
  450. tm.assert_index_equal(result, expected)
  451. i2 = i.copy()
  452. i2 = PeriodIndex([NaT, NaT] + i[2:].tolist(), freq="D")
  453. result = i.where(notna(i2), i2)
  454. tm.assert_index_equal(result, i2)
  455. i2 = i.copy()
  456. i2 = PeriodIndex([NaT, NaT] + i[2:].tolist(), freq="D")
  457. result = i.where(notna(i2), i2.values)
  458. tm.assert_index_equal(result, i2)
  459. def test_where_invalid_dtypes(self):
  460. pi = period_range("20130101", periods=5, freq="D")
  461. tail = pi[2:].tolist()
  462. i2 = PeriodIndex([NaT, NaT] + tail, freq="D")
  463. mask = notna(i2)
  464. result = pi.where(mask, i2.asi8)
  465. expected = pd.Index([NaT._value, NaT._value] + tail, dtype=object)
  466. assert isinstance(expected[0], int)
  467. tm.assert_index_equal(result, expected)
  468. tdi = i2.asi8.view("timedelta64[ns]")
  469. expected = pd.Index([tdi[0], tdi[1]] + tail, dtype=object)
  470. assert isinstance(expected[0], np.timedelta64)
  471. result = pi.where(mask, tdi)
  472. tm.assert_index_equal(result, expected)
  473. dti = i2.to_timestamp("S")
  474. expected = pd.Index([dti[0], dti[1]] + tail, dtype=object)
  475. assert expected[0] is NaT
  476. result = pi.where(mask, dti)
  477. tm.assert_index_equal(result, expected)
  478. td = Timedelta(days=4)
  479. expected = pd.Index([td, td] + tail, dtype=object)
  480. assert expected[0] == td
  481. result = pi.where(mask, td)
  482. tm.assert_index_equal(result, expected)
  483. def test_where_mismatched_nat(self):
  484. pi = period_range("20130101", periods=5, freq="D")
  485. cond = np.array([True, False, True, True, False])
  486. tdnat = np.timedelta64("NaT", "ns")
  487. expected = pd.Index([pi[0], tdnat, pi[2], pi[3], tdnat], dtype=object)
  488. assert expected[1] is tdnat
  489. result = pi.where(cond, tdnat)
  490. tm.assert_index_equal(result, expected)
  491. class TestTake:
  492. def test_take(self):
  493. # GH#10295
  494. idx1 = period_range("2011-01-01", "2011-01-31", freq="D", name="idx")
  495. for idx in [idx1]:
  496. result = idx.take([0])
  497. assert result == Period("2011-01-01", freq="D")
  498. result = idx.take([5])
  499. assert result == Period("2011-01-06", freq="D")
  500. result = idx.take([0, 1, 2])
  501. expected = period_range("2011-01-01", "2011-01-03", freq="D", name="idx")
  502. tm.assert_index_equal(result, expected)
  503. assert result.freq == "D"
  504. assert result.freq == expected.freq
  505. result = idx.take([0, 2, 4])
  506. expected = PeriodIndex(
  507. ["2011-01-01", "2011-01-03", "2011-01-05"], freq="D", name="idx"
  508. )
  509. tm.assert_index_equal(result, expected)
  510. assert result.freq == expected.freq
  511. assert result.freq == "D"
  512. result = idx.take([7, 4, 1])
  513. expected = PeriodIndex(
  514. ["2011-01-08", "2011-01-05", "2011-01-02"], freq="D", name="idx"
  515. )
  516. tm.assert_index_equal(result, expected)
  517. assert result.freq == expected.freq
  518. assert result.freq == "D"
  519. result = idx.take([3, 2, 5])
  520. expected = PeriodIndex(
  521. ["2011-01-04", "2011-01-03", "2011-01-06"], freq="D", name="idx"
  522. )
  523. tm.assert_index_equal(result, expected)
  524. assert result.freq == expected.freq
  525. assert result.freq == "D"
  526. result = idx.take([-3, 2, 5])
  527. expected = PeriodIndex(
  528. ["2011-01-29", "2011-01-03", "2011-01-06"], freq="D", name="idx"
  529. )
  530. tm.assert_index_equal(result, expected)
  531. assert result.freq == expected.freq
  532. assert result.freq == "D"
  533. def test_take_misc(self):
  534. index = period_range(start="1/1/10", end="12/31/12", freq="D", name="idx")
  535. expected = PeriodIndex(
  536. [
  537. datetime(2010, 1, 6),
  538. datetime(2010, 1, 7),
  539. datetime(2010, 1, 9),
  540. datetime(2010, 1, 13),
  541. ],
  542. freq="D",
  543. name="idx",
  544. )
  545. taken1 = index.take([5, 6, 8, 12])
  546. taken2 = index[[5, 6, 8, 12]]
  547. for taken in [taken1, taken2]:
  548. tm.assert_index_equal(taken, expected)
  549. assert isinstance(taken, PeriodIndex)
  550. assert taken.freq == index.freq
  551. assert taken.name == expected.name
  552. def test_take_fill_value(self):
  553. # GH#12631
  554. idx = PeriodIndex(
  555. ["2011-01-01", "2011-02-01", "2011-03-01"], name="xxx", freq="D"
  556. )
  557. result = idx.take(np.array([1, 0, -1]))
  558. expected = PeriodIndex(
  559. ["2011-02-01", "2011-01-01", "2011-03-01"], name="xxx", freq="D"
  560. )
  561. tm.assert_index_equal(result, expected)
  562. # fill_value
  563. result = idx.take(np.array([1, 0, -1]), fill_value=True)
  564. expected = PeriodIndex(
  565. ["2011-02-01", "2011-01-01", "NaT"], name="xxx", freq="D"
  566. )
  567. tm.assert_index_equal(result, expected)
  568. # allow_fill=False
  569. result = idx.take(np.array([1, 0, -1]), allow_fill=False, fill_value=True)
  570. expected = PeriodIndex(
  571. ["2011-02-01", "2011-01-01", "2011-03-01"], name="xxx", freq="D"
  572. )
  573. tm.assert_index_equal(result, expected)
  574. msg = (
  575. "When allow_fill=True and fill_value is not None, "
  576. "all indices must be >= -1"
  577. )
  578. with pytest.raises(ValueError, match=msg):
  579. idx.take(np.array([1, 0, -2]), fill_value=True)
  580. with pytest.raises(ValueError, match=msg):
  581. idx.take(np.array([1, 0, -5]), fill_value=True)
  582. msg = "index -5 is out of bounds for( axis 0 with)? size 3"
  583. with pytest.raises(IndexError, match=msg):
  584. idx.take(np.array([1, -5]))
  585. class TestGetValue:
  586. @pytest.mark.parametrize("freq", ["H", "D"])
  587. def test_get_value_datetime_hourly(self, freq):
  588. # get_loc and get_value should treat datetime objects symmetrically
  589. # TODO: this test used to test get_value, which is removed in 2.0.
  590. # should this test be moved somewhere, or is what's left redundant?
  591. dti = date_range("2016-01-01", periods=3, freq="MS")
  592. pi = dti.to_period(freq)
  593. ser = Series(range(7, 10), index=pi)
  594. ts = dti[0]
  595. assert pi.get_loc(ts) == 0
  596. assert ser[ts] == 7
  597. assert ser.loc[ts] == 7
  598. ts2 = ts + Timedelta(hours=3)
  599. if freq == "H":
  600. with pytest.raises(KeyError, match="2016-01-01 03:00"):
  601. pi.get_loc(ts2)
  602. with pytest.raises(KeyError, match="2016-01-01 03:00"):
  603. ser[ts2]
  604. with pytest.raises(KeyError, match="2016-01-01 03:00"):
  605. ser.loc[ts2]
  606. else:
  607. assert pi.get_loc(ts2) == 0
  608. assert ser[ts2] == 7
  609. assert ser.loc[ts2] == 7
  610. class TestContains:
  611. def test_contains(self):
  612. # GH 17717
  613. p0 = Period("2017-09-01")
  614. p1 = Period("2017-09-02")
  615. p2 = Period("2017-09-03")
  616. p3 = Period("2017-09-04")
  617. ps0 = [p0, p1, p2]
  618. idx0 = PeriodIndex(ps0)
  619. for p in ps0:
  620. assert p in idx0
  621. assert str(p) in idx0
  622. # GH#31172
  623. # Higher-resolution period-like are _not_ considered as contained
  624. key = "2017-09-01 00:00:01"
  625. assert key not in idx0
  626. with pytest.raises(KeyError, match=key):
  627. idx0.get_loc(key)
  628. assert "2017-09" in idx0
  629. assert p3 not in idx0
  630. def test_contains_freq_mismatch(self):
  631. rng = period_range("2007-01", freq="M", periods=10)
  632. assert Period("2007-01", freq="M") in rng
  633. assert Period("2007-01", freq="D") not in rng
  634. assert Period("2007-01", freq="2M") not in rng
  635. def test_contains_nat(self):
  636. # see gh-13582
  637. idx = period_range("2007-01", freq="M", periods=10)
  638. assert NaT not in idx
  639. assert None not in idx
  640. assert float("nan") not in idx
  641. assert np.nan not in idx
  642. idx = PeriodIndex(["2011-01", "NaT", "2011-02"], freq="M")
  643. assert NaT in idx
  644. assert None in idx
  645. assert float("nan") in idx
  646. assert np.nan in idx
  647. class TestAsOfLocs:
  648. def test_asof_locs_mismatched_type(self):
  649. dti = date_range("2016-01-01", periods=3)
  650. pi = dti.to_period("D")
  651. pi2 = dti.to_period("H")
  652. mask = np.array([0, 1, 0], dtype=bool)
  653. msg = "must be DatetimeIndex or PeriodIndex"
  654. with pytest.raises(TypeError, match=msg):
  655. pi.asof_locs(pd.Index(pi.asi8, dtype=np.int64), mask)
  656. with pytest.raises(TypeError, match=msg):
  657. pi.asof_locs(pd.Index(pi.asi8, dtype=np.float64), mask)
  658. with pytest.raises(TypeError, match=msg):
  659. # TimedeltaIndex
  660. pi.asof_locs(dti - dti, mask)
  661. msg = "Input has different freq=H"
  662. with pytest.raises(libperiod.IncompatibleFrequency, match=msg):
  663. pi.asof_locs(pi2, mask)