test_partial.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668
  1. """
  2. test setting *parts* of objects both positionally and label based
  3. TODO: these should be split among the indexer tests
  4. """
  5. import numpy as np
  6. import pytest
  7. import pandas as pd
  8. from pandas import (
  9. DataFrame,
  10. Index,
  11. Period,
  12. Series,
  13. Timestamp,
  14. date_range,
  15. period_range,
  16. )
  17. import pandas._testing as tm
  18. class TestEmptyFrameSetitemExpansion:
  19. def test_empty_frame_setitem_index_name_retained(self):
  20. # GH#31368 empty frame has non-None index.name -> retained
  21. df = DataFrame({}, index=pd.RangeIndex(0, name="df_index"))
  22. series = Series(1.23, index=pd.RangeIndex(4, name="series_index"))
  23. df["series"] = series
  24. expected = DataFrame(
  25. {"series": [1.23] * 4}, index=pd.RangeIndex(4, name="df_index")
  26. )
  27. tm.assert_frame_equal(df, expected)
  28. def test_empty_frame_setitem_index_name_inherited(self):
  29. # GH#36527 empty frame has None index.name -> not retained
  30. df = DataFrame()
  31. series = Series(1.23, index=pd.RangeIndex(4, name="series_index"))
  32. df["series"] = series
  33. expected = DataFrame(
  34. {"series": [1.23] * 4}, index=pd.RangeIndex(4, name="series_index")
  35. )
  36. tm.assert_frame_equal(df, expected)
  37. def test_loc_setitem_zerolen_series_columns_align(self):
  38. # columns will align
  39. df = DataFrame(columns=["A", "B"])
  40. df.loc[0] = Series(1, index=range(4))
  41. expected = DataFrame(columns=["A", "B"], index=[0], dtype=np.float64)
  42. tm.assert_frame_equal(df, expected)
  43. # columns will align
  44. df = DataFrame(columns=["A", "B"])
  45. df.loc[0] = Series(1, index=["B"])
  46. exp = DataFrame([[np.nan, 1]], columns=["A", "B"], index=[0], dtype="float64")
  47. tm.assert_frame_equal(df, exp)
  48. def test_loc_setitem_zerolen_list_length_must_match_columns(self):
  49. # list-like must conform
  50. df = DataFrame(columns=["A", "B"])
  51. msg = "cannot set a row with mismatched columns"
  52. with pytest.raises(ValueError, match=msg):
  53. df.loc[0] = [1, 2, 3]
  54. df = DataFrame(columns=["A", "B"])
  55. df.loc[3] = [6, 7] # length matches len(df.columns) --> OK!
  56. exp = DataFrame([[6, 7]], index=[3], columns=["A", "B"], dtype=np.int64)
  57. tm.assert_frame_equal(df, exp)
  58. def test_partial_set_empty_frame(self):
  59. # partially set with an empty object
  60. # frame
  61. df = DataFrame()
  62. msg = "cannot set a frame with no defined columns"
  63. with pytest.raises(ValueError, match=msg):
  64. df.loc[1] = 1
  65. with pytest.raises(ValueError, match=msg):
  66. df.loc[1] = Series([1], index=["foo"])
  67. msg = "cannot set a frame with no defined index and a scalar"
  68. with pytest.raises(ValueError, match=msg):
  69. df.loc[:, 1] = 1
  70. def test_partial_set_empty_frame2(self):
  71. # these work as they don't really change
  72. # anything but the index
  73. # GH#5632
  74. expected = DataFrame(columns=["foo"], index=Index([], dtype="object"))
  75. df = DataFrame(index=Index([], dtype="object"))
  76. df["foo"] = Series([], dtype="object")
  77. tm.assert_frame_equal(df, expected)
  78. df = DataFrame(index=Index([]))
  79. df["foo"] = Series(df.index)
  80. tm.assert_frame_equal(df, expected)
  81. df = DataFrame(index=Index([]))
  82. df["foo"] = df.index
  83. tm.assert_frame_equal(df, expected)
  84. def test_partial_set_empty_frame3(self):
  85. expected = DataFrame(columns=["foo"], index=Index([], dtype="int64"))
  86. expected["foo"] = expected["foo"].astype("float64")
  87. df = DataFrame(index=Index([], dtype="int64"))
  88. df["foo"] = []
  89. tm.assert_frame_equal(df, expected)
  90. df = DataFrame(index=Index([], dtype="int64"))
  91. df["foo"] = Series(np.arange(len(df)), dtype="float64")
  92. tm.assert_frame_equal(df, expected)
  93. def test_partial_set_empty_frame4(self):
  94. df = DataFrame(index=Index([], dtype="int64"))
  95. df["foo"] = range(len(df))
  96. expected = DataFrame(columns=["foo"], index=Index([], dtype="int64"))
  97. # range is int-dtype-like, so we get int64 dtype
  98. expected["foo"] = expected["foo"].astype("int64")
  99. tm.assert_frame_equal(df, expected)
  100. def test_partial_set_empty_frame5(self):
  101. df = DataFrame()
  102. tm.assert_index_equal(df.columns, pd.RangeIndex(0))
  103. df2 = DataFrame()
  104. df2[1] = Series([1], index=["foo"])
  105. df.loc[:, 1] = Series([1], index=["foo"])
  106. tm.assert_frame_equal(df, DataFrame([[1]], index=["foo"], columns=[1]))
  107. tm.assert_frame_equal(df, df2)
  108. def test_partial_set_empty_frame_no_index(self):
  109. # no index to start
  110. expected = DataFrame({0: Series(1, index=range(4))}, columns=["A", "B", 0])
  111. df = DataFrame(columns=["A", "B"])
  112. df[0] = Series(1, index=range(4))
  113. df.dtypes
  114. str(df)
  115. tm.assert_frame_equal(df, expected)
  116. df = DataFrame(columns=["A", "B"])
  117. df.loc[:, 0] = Series(1, index=range(4))
  118. df.dtypes
  119. str(df)
  120. tm.assert_frame_equal(df, expected)
  121. def test_partial_set_empty_frame_row(self):
  122. # GH#5720, GH#5744
  123. # don't create rows when empty
  124. expected = DataFrame(columns=["A", "B", "New"], index=Index([], dtype="int64"))
  125. expected["A"] = expected["A"].astype("int64")
  126. expected["B"] = expected["B"].astype("float64")
  127. expected["New"] = expected["New"].astype("float64")
  128. df = DataFrame({"A": [1, 2, 3], "B": [1.2, 4.2, 5.2]})
  129. y = df[df.A > 5]
  130. y["New"] = np.nan
  131. tm.assert_frame_equal(y, expected)
  132. expected = DataFrame(columns=["a", "b", "c c", "d"])
  133. expected["d"] = expected["d"].astype("int64")
  134. df = DataFrame(columns=["a", "b", "c c"])
  135. df["d"] = 3
  136. tm.assert_frame_equal(df, expected)
  137. tm.assert_series_equal(df["c c"], Series(name="c c", dtype=object))
  138. # reindex columns is ok
  139. df = DataFrame({"A": [1, 2, 3], "B": [1.2, 4.2, 5.2]})
  140. y = df[df.A > 5]
  141. result = y.reindex(columns=["A", "B", "C"])
  142. expected = DataFrame(columns=["A", "B", "C"])
  143. expected["A"] = expected["A"].astype("int64")
  144. expected["B"] = expected["B"].astype("float64")
  145. expected["C"] = expected["C"].astype("float64")
  146. tm.assert_frame_equal(result, expected)
  147. def test_partial_set_empty_frame_set_series(self):
  148. # GH#5756
  149. # setting with empty Series
  150. df = DataFrame(Series(dtype=object))
  151. expected = DataFrame({0: Series(dtype=object)})
  152. tm.assert_frame_equal(df, expected)
  153. df = DataFrame(Series(name="foo", dtype=object))
  154. expected = DataFrame({"foo": Series(dtype=object)})
  155. tm.assert_frame_equal(df, expected)
  156. def test_partial_set_empty_frame_empty_copy_assignment(self):
  157. # GH#5932
  158. # copy on empty with assignment fails
  159. df = DataFrame(index=[0])
  160. df = df.copy()
  161. df["a"] = 0
  162. expected = DataFrame(0, index=[0], columns=["a"])
  163. tm.assert_frame_equal(df, expected)
  164. def test_partial_set_empty_frame_empty_consistencies(self):
  165. # GH#6171
  166. # consistency on empty frames
  167. df = DataFrame(columns=["x", "y"])
  168. df["x"] = [1, 2]
  169. expected = DataFrame({"x": [1, 2], "y": [np.nan, np.nan]})
  170. tm.assert_frame_equal(df, expected, check_dtype=False)
  171. df = DataFrame(columns=["x", "y"])
  172. df["x"] = ["1", "2"]
  173. expected = DataFrame({"x": ["1", "2"], "y": [np.nan, np.nan]}, dtype=object)
  174. tm.assert_frame_equal(df, expected)
  175. df = DataFrame(columns=["x", "y"])
  176. df.loc[0, "x"] = 1
  177. expected = DataFrame({"x": [1], "y": [np.nan]})
  178. tm.assert_frame_equal(df, expected, check_dtype=False)
  179. class TestPartialSetting:
  180. def test_partial_setting(self):
  181. # GH2578, allow ix and friends to partially set
  182. # series
  183. s_orig = Series([1, 2, 3])
  184. s = s_orig.copy()
  185. s[5] = 5
  186. expected = Series([1, 2, 3, 5], index=[0, 1, 2, 5])
  187. tm.assert_series_equal(s, expected)
  188. s = s_orig.copy()
  189. s.loc[5] = 5
  190. expected = Series([1, 2, 3, 5], index=[0, 1, 2, 5])
  191. tm.assert_series_equal(s, expected)
  192. s = s_orig.copy()
  193. s[5] = 5.0
  194. expected = Series([1, 2, 3, 5.0], index=[0, 1, 2, 5])
  195. tm.assert_series_equal(s, expected)
  196. s = s_orig.copy()
  197. s.loc[5] = 5.0
  198. expected = Series([1, 2, 3, 5.0], index=[0, 1, 2, 5])
  199. tm.assert_series_equal(s, expected)
  200. # iloc/iat raise
  201. s = s_orig.copy()
  202. msg = "iloc cannot enlarge its target object"
  203. with pytest.raises(IndexError, match=msg):
  204. s.iloc[3] = 5.0
  205. msg = "index 3 is out of bounds for axis 0 with size 3"
  206. with pytest.raises(IndexError, match=msg):
  207. s.iat[3] = 5.0
  208. def test_partial_setting_frame(self, using_array_manager):
  209. df_orig = DataFrame(
  210. np.arange(6).reshape(3, 2), columns=["A", "B"], dtype="int64"
  211. )
  212. # iloc/iat raise
  213. df = df_orig.copy()
  214. msg = "iloc cannot enlarge its target object"
  215. with pytest.raises(IndexError, match=msg):
  216. df.iloc[4, 2] = 5.0
  217. msg = "index 2 is out of bounds for axis 0 with size 2"
  218. if using_array_manager:
  219. msg = "list index out of range"
  220. with pytest.raises(IndexError, match=msg):
  221. df.iat[4, 2] = 5.0
  222. # row setting where it exists
  223. expected = DataFrame(dict({"A": [0, 4, 4], "B": [1, 5, 5]}))
  224. df = df_orig.copy()
  225. df.iloc[1] = df.iloc[2]
  226. tm.assert_frame_equal(df, expected)
  227. expected = DataFrame(dict({"A": [0, 4, 4], "B": [1, 5, 5]}))
  228. df = df_orig.copy()
  229. df.loc[1] = df.loc[2]
  230. tm.assert_frame_equal(df, expected)
  231. # like 2578, partial setting with dtype preservation
  232. expected = DataFrame(dict({"A": [0, 2, 4, 4], "B": [1, 3, 5, 5]}))
  233. df = df_orig.copy()
  234. df.loc[3] = df.loc[2]
  235. tm.assert_frame_equal(df, expected)
  236. # single dtype frame, overwrite
  237. expected = DataFrame(dict({"A": [0, 2, 4], "B": [0, 2, 4]}))
  238. df = df_orig.copy()
  239. df.loc[:, "B"] = df.loc[:, "A"]
  240. tm.assert_frame_equal(df, expected)
  241. # mixed dtype frame, overwrite
  242. expected = DataFrame(dict({"A": [0, 2, 4], "B": Series([0.0, 2.0, 4.0])}))
  243. df = df_orig.copy()
  244. df["B"] = df["B"].astype(np.float64)
  245. # as of 2.0, df.loc[:, "B"] = ... attempts (and here succeeds) at
  246. # setting inplace
  247. df.loc[:, "B"] = df.loc[:, "A"]
  248. tm.assert_frame_equal(df, expected)
  249. # single dtype frame, partial setting
  250. expected = df_orig.copy()
  251. expected["C"] = df["A"]
  252. df = df_orig.copy()
  253. df.loc[:, "C"] = df.loc[:, "A"]
  254. tm.assert_frame_equal(df, expected)
  255. # mixed frame, partial setting
  256. expected = df_orig.copy()
  257. expected["C"] = df["A"]
  258. df = df_orig.copy()
  259. df.loc[:, "C"] = df.loc[:, "A"]
  260. tm.assert_frame_equal(df, expected)
  261. def test_partial_setting2(self):
  262. # GH 8473
  263. dates = date_range("1/1/2000", periods=8)
  264. df_orig = DataFrame(
  265. np.random.randn(8, 4), index=dates, columns=["A", "B", "C", "D"]
  266. )
  267. expected = pd.concat(
  268. [df_orig, DataFrame({"A": 7}, index=dates[-1:] + dates.freq)], sort=True
  269. )
  270. df = df_orig.copy()
  271. df.loc[dates[-1] + dates.freq, "A"] = 7
  272. tm.assert_frame_equal(df, expected)
  273. df = df_orig.copy()
  274. df.at[dates[-1] + dates.freq, "A"] = 7
  275. tm.assert_frame_equal(df, expected)
  276. exp_other = DataFrame({0: 7}, index=dates[-1:] + dates.freq)
  277. expected = pd.concat([df_orig, exp_other], axis=1)
  278. df = df_orig.copy()
  279. df.loc[dates[-1] + dates.freq, 0] = 7
  280. tm.assert_frame_equal(df, expected)
  281. df = df_orig.copy()
  282. df.at[dates[-1] + dates.freq, 0] = 7
  283. tm.assert_frame_equal(df, expected)
  284. def test_partial_setting_mixed_dtype(self):
  285. # in a mixed dtype environment, try to preserve dtypes
  286. # by appending
  287. df = DataFrame([[True, 1], [False, 2]], columns=["female", "fitness"])
  288. s = df.loc[1].copy()
  289. s.name = 2
  290. expected = pd.concat([df, DataFrame(s).T.infer_objects()])
  291. df.loc[2] = df.loc[1]
  292. tm.assert_frame_equal(df, expected)
  293. def test_series_partial_set(self):
  294. # partial set with new index
  295. # Regression from GH4825
  296. ser = Series([0.1, 0.2], index=[1, 2])
  297. # loc equiv to .reindex
  298. expected = Series([np.nan, 0.2, np.nan], index=[3, 2, 3])
  299. with pytest.raises(KeyError, match=r"not in index"):
  300. ser.loc[[3, 2, 3]]
  301. result = ser.reindex([3, 2, 3])
  302. tm.assert_series_equal(result, expected, check_index_type=True)
  303. expected = Series([np.nan, 0.2, np.nan, np.nan], index=[3, 2, 3, "x"])
  304. with pytest.raises(KeyError, match="not in index"):
  305. ser.loc[[3, 2, 3, "x"]]
  306. result = ser.reindex([3, 2, 3, "x"])
  307. tm.assert_series_equal(result, expected, check_index_type=True)
  308. expected = Series([0.2, 0.2, 0.1], index=[2, 2, 1])
  309. result = ser.loc[[2, 2, 1]]
  310. tm.assert_series_equal(result, expected, check_index_type=True)
  311. expected = Series([0.2, 0.2, np.nan, 0.1], index=[2, 2, "x", 1])
  312. with pytest.raises(KeyError, match="not in index"):
  313. ser.loc[[2, 2, "x", 1]]
  314. result = ser.reindex([2, 2, "x", 1])
  315. tm.assert_series_equal(result, expected, check_index_type=True)
  316. # raises as nothing is in the index
  317. msg = (
  318. rf"\"None of \[Index\(\[3, 3, 3\], dtype='{np.int_().dtype}'\)\] "
  319. r"are in the \[index\]\""
  320. )
  321. with pytest.raises(KeyError, match=msg):
  322. ser.loc[[3, 3, 3]]
  323. expected = Series([0.2, 0.2, np.nan], index=[2, 2, 3])
  324. with pytest.raises(KeyError, match="not in index"):
  325. ser.loc[[2, 2, 3]]
  326. result = ser.reindex([2, 2, 3])
  327. tm.assert_series_equal(result, expected, check_index_type=True)
  328. s = Series([0.1, 0.2, 0.3], index=[1, 2, 3])
  329. expected = Series([0.3, np.nan, np.nan], index=[3, 4, 4])
  330. with pytest.raises(KeyError, match="not in index"):
  331. s.loc[[3, 4, 4]]
  332. result = s.reindex([3, 4, 4])
  333. tm.assert_series_equal(result, expected, check_index_type=True)
  334. s = Series([0.1, 0.2, 0.3, 0.4], index=[1, 2, 3, 4])
  335. expected = Series([np.nan, 0.3, 0.3], index=[5, 3, 3])
  336. with pytest.raises(KeyError, match="not in index"):
  337. s.loc[[5, 3, 3]]
  338. result = s.reindex([5, 3, 3])
  339. tm.assert_series_equal(result, expected, check_index_type=True)
  340. s = Series([0.1, 0.2, 0.3, 0.4], index=[1, 2, 3, 4])
  341. expected = Series([np.nan, 0.4, 0.4], index=[5, 4, 4])
  342. with pytest.raises(KeyError, match="not in index"):
  343. s.loc[[5, 4, 4]]
  344. result = s.reindex([5, 4, 4])
  345. tm.assert_series_equal(result, expected, check_index_type=True)
  346. s = Series([0.1, 0.2, 0.3, 0.4], index=[4, 5, 6, 7])
  347. expected = Series([0.4, np.nan, np.nan], index=[7, 2, 2])
  348. with pytest.raises(KeyError, match="not in index"):
  349. s.loc[[7, 2, 2]]
  350. result = s.reindex([7, 2, 2])
  351. tm.assert_series_equal(result, expected, check_index_type=True)
  352. s = Series([0.1, 0.2, 0.3, 0.4], index=[1, 2, 3, 4])
  353. expected = Series([0.4, np.nan, np.nan], index=[4, 5, 5])
  354. with pytest.raises(KeyError, match="not in index"):
  355. s.loc[[4, 5, 5]]
  356. result = s.reindex([4, 5, 5])
  357. tm.assert_series_equal(result, expected, check_index_type=True)
  358. # iloc
  359. expected = Series([0.2, 0.2, 0.1, 0.1], index=[2, 2, 1, 1])
  360. result = ser.iloc[[1, 1, 0, 0]]
  361. tm.assert_series_equal(result, expected, check_index_type=True)
  362. def test_series_partial_set_with_name(self):
  363. # GH 11497
  364. idx = Index([1, 2], dtype="int64", name="idx")
  365. ser = Series([0.1, 0.2], index=idx, name="s")
  366. # loc
  367. with pytest.raises(KeyError, match=r"\[3\] not in index"):
  368. ser.loc[[3, 2, 3]]
  369. with pytest.raises(KeyError, match=r"not in index"):
  370. ser.loc[[3, 2, 3, "x"]]
  371. exp_idx = Index([2, 2, 1], dtype="int64", name="idx")
  372. expected = Series([0.2, 0.2, 0.1], index=exp_idx, name="s")
  373. result = ser.loc[[2, 2, 1]]
  374. tm.assert_series_equal(result, expected, check_index_type=True)
  375. with pytest.raises(KeyError, match=r"\['x'\] not in index"):
  376. ser.loc[[2, 2, "x", 1]]
  377. # raises as nothing is in the index
  378. msg = (
  379. rf"\"None of \[Index\(\[3, 3, 3\], dtype='{np.int_().dtype}', "
  380. r"name='idx'\)\] are in the \[index\]\""
  381. )
  382. with pytest.raises(KeyError, match=msg):
  383. ser.loc[[3, 3, 3]]
  384. with pytest.raises(KeyError, match="not in index"):
  385. ser.loc[[2, 2, 3]]
  386. idx = Index([1, 2, 3], dtype="int64", name="idx")
  387. with pytest.raises(KeyError, match="not in index"):
  388. Series([0.1, 0.2, 0.3], index=idx, name="s").loc[[3, 4, 4]]
  389. idx = Index([1, 2, 3, 4], dtype="int64", name="idx")
  390. with pytest.raises(KeyError, match="not in index"):
  391. Series([0.1, 0.2, 0.3, 0.4], index=idx, name="s").loc[[5, 3, 3]]
  392. idx = Index([1, 2, 3, 4], dtype="int64", name="idx")
  393. with pytest.raises(KeyError, match="not in index"):
  394. Series([0.1, 0.2, 0.3, 0.4], index=idx, name="s").loc[[5, 4, 4]]
  395. idx = Index([4, 5, 6, 7], dtype="int64", name="idx")
  396. with pytest.raises(KeyError, match="not in index"):
  397. Series([0.1, 0.2, 0.3, 0.4], index=idx, name="s").loc[[7, 2, 2]]
  398. idx = Index([1, 2, 3, 4], dtype="int64", name="idx")
  399. with pytest.raises(KeyError, match="not in index"):
  400. Series([0.1, 0.2, 0.3, 0.4], index=idx, name="s").loc[[4, 5, 5]]
  401. # iloc
  402. exp_idx = Index([2, 2, 1, 1], dtype="int64", name="idx")
  403. expected = Series([0.2, 0.2, 0.1, 0.1], index=exp_idx, name="s")
  404. result = ser.iloc[[1, 1, 0, 0]]
  405. tm.assert_series_equal(result, expected, check_index_type=True)
  406. @pytest.mark.parametrize("key", [100, 100.0])
  407. def test_setitem_with_expansion_numeric_into_datetimeindex(self, key):
  408. # GH#4940 inserting non-strings
  409. orig = tm.makeTimeDataFrame()
  410. df = orig.copy()
  411. df.loc[key, :] = df.iloc[0]
  412. ex_index = Index(list(orig.index) + [key], dtype=object, name=orig.index.name)
  413. ex_data = np.concatenate([orig.values, df.iloc[[0]].values], axis=0)
  414. expected = DataFrame(ex_data, index=ex_index, columns=orig.columns)
  415. tm.assert_frame_equal(df, expected)
  416. def test_partial_set_invalid(self):
  417. # GH 4940
  418. # allow only setting of 'valid' values
  419. orig = tm.makeTimeDataFrame()
  420. # allow object conversion here
  421. df = orig.copy()
  422. df.loc["a", :] = df.iloc[0]
  423. ser = Series(df.iloc[0], name="a")
  424. exp = pd.concat([orig, DataFrame(ser).T.infer_objects()])
  425. tm.assert_frame_equal(df, exp)
  426. tm.assert_index_equal(df.index, Index(orig.index.tolist() + ["a"]))
  427. assert df.index.dtype == "object"
  428. @pytest.mark.parametrize(
  429. "idx,labels,expected_idx",
  430. [
  431. (
  432. period_range(start="2000", periods=20, freq="D"),
  433. ["2000-01-04", "2000-01-08", "2000-01-12"],
  434. [
  435. Period("2000-01-04", freq="D"),
  436. Period("2000-01-08", freq="D"),
  437. Period("2000-01-12", freq="D"),
  438. ],
  439. ),
  440. (
  441. date_range(start="2000", periods=20, freq="D"),
  442. ["2000-01-04", "2000-01-08", "2000-01-12"],
  443. [
  444. Timestamp("2000-01-04"),
  445. Timestamp("2000-01-08"),
  446. Timestamp("2000-01-12"),
  447. ],
  448. ),
  449. (
  450. pd.timedelta_range(start="1 day", periods=20),
  451. ["4D", "8D", "12D"],
  452. [pd.Timedelta("4 day"), pd.Timedelta("8 day"), pd.Timedelta("12 day")],
  453. ),
  454. ],
  455. )
  456. def test_loc_with_list_of_strings_representing_datetimes(
  457. self, idx, labels, expected_idx, frame_or_series
  458. ):
  459. # GH 11278
  460. obj = frame_or_series(range(20), index=idx)
  461. expected_value = [3, 7, 11]
  462. expected = frame_or_series(expected_value, expected_idx)
  463. tm.assert_equal(expected, obj.loc[labels])
  464. if frame_or_series is Series:
  465. tm.assert_series_equal(expected, obj[labels])
  466. @pytest.mark.parametrize(
  467. "idx,labels",
  468. [
  469. (
  470. period_range(start="2000", periods=20, freq="D"),
  471. ["2000-01-04", "2000-01-30"],
  472. ),
  473. (
  474. date_range(start="2000", periods=20, freq="D"),
  475. ["2000-01-04", "2000-01-30"],
  476. ),
  477. (pd.timedelta_range(start="1 day", periods=20), ["3 day", "30 day"]),
  478. ],
  479. )
  480. def test_loc_with_list_of_strings_representing_datetimes_missing_value(
  481. self, idx, labels
  482. ):
  483. # GH 11278
  484. ser = Series(range(20), index=idx)
  485. df = DataFrame(range(20), index=idx)
  486. msg = r"not in index"
  487. with pytest.raises(KeyError, match=msg):
  488. ser.loc[labels]
  489. with pytest.raises(KeyError, match=msg):
  490. ser[labels]
  491. with pytest.raises(KeyError, match=msg):
  492. df.loc[labels]
  493. @pytest.mark.parametrize(
  494. "idx,labels,msg",
  495. [
  496. (
  497. period_range(start="2000", periods=20, freq="D"),
  498. ["4D", "8D"],
  499. (
  500. r"None of \[Index\(\['4D', '8D'\], dtype='object'\)\] "
  501. r"are in the \[index\]"
  502. ),
  503. ),
  504. (
  505. date_range(start="2000", periods=20, freq="D"),
  506. ["4D", "8D"],
  507. (
  508. r"None of \[Index\(\['4D', '8D'\], dtype='object'\)\] "
  509. r"are in the \[index\]"
  510. ),
  511. ),
  512. (
  513. pd.timedelta_range(start="1 day", periods=20),
  514. ["2000-01-04", "2000-01-08"],
  515. (
  516. r"None of \[Index\(\['2000-01-04', '2000-01-08'\], "
  517. r"dtype='object'\)\] are in the \[index\]"
  518. ),
  519. ),
  520. ],
  521. )
  522. def test_loc_with_list_of_strings_representing_datetimes_not_matched_type(
  523. self, idx, labels, msg
  524. ):
  525. # GH 11278
  526. ser = Series(range(20), index=idx)
  527. df = DataFrame(range(20), index=idx)
  528. with pytest.raises(KeyError, match=msg):
  529. ser.loc[labels]
  530. with pytest.raises(KeyError, match=msg):
  531. ser[labels]
  532. with pytest.raises(KeyError, match=msg):
  533. df.loc[labels]
  534. class TestStringSlicing:
  535. def test_slice_irregular_datetime_index_with_nan(self):
  536. # GH36953
  537. index = pd.to_datetime(["2012-01-01", "2012-01-02", "2012-01-03", None])
  538. df = DataFrame(range(len(index)), index=index)
  539. expected = DataFrame(range(len(index[:3])), index=index[:3])
  540. result = df["2012-01-01":"2012-01-04"]
  541. tm.assert_frame_equal(result, expected)