test_scalar.py 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297
  1. """ test scalar indexing, including at and iat """
  2. from datetime import (
  3. datetime,
  4. timedelta,
  5. )
  6. import itertools
  7. import numpy as np
  8. import pytest
  9. from pandas import (
  10. DataFrame,
  11. Series,
  12. Timedelta,
  13. Timestamp,
  14. date_range,
  15. )
  16. import pandas._testing as tm
  17. def generate_indices(f, values=False):
  18. """
  19. generate the indices
  20. if values is True , use the axis values
  21. is False, use the range
  22. """
  23. axes = f.axes
  24. if values:
  25. axes = (list(range(len(ax))) for ax in axes)
  26. return itertools.product(*axes)
  27. class TestScalar:
  28. @pytest.mark.parametrize("kind", ["series", "frame"])
  29. @pytest.mark.parametrize("col", ["ints", "uints"])
  30. def test_iat_set_ints(self, kind, col, request):
  31. f = request.getfixturevalue(f"{kind}_{col}")
  32. indices = generate_indices(f, True)
  33. for i in indices:
  34. f.iat[i] = 1
  35. expected = f.values[i]
  36. tm.assert_almost_equal(expected, 1)
  37. @pytest.mark.parametrize("kind", ["series", "frame"])
  38. @pytest.mark.parametrize("col", ["labels", "ts", "floats"])
  39. def test_iat_set_other(self, kind, col, request):
  40. f = request.getfixturevalue(f"{kind}_{col}")
  41. msg = "iAt based indexing can only have integer indexers"
  42. with pytest.raises(ValueError, match=msg):
  43. idx = next(generate_indices(f, False))
  44. f.iat[idx] = 1
  45. @pytest.mark.parametrize("kind", ["series", "frame"])
  46. @pytest.mark.parametrize("col", ["ints", "uints", "labels", "ts", "floats"])
  47. def test_at_set_ints_other(self, kind, col, request):
  48. f = request.getfixturevalue(f"{kind}_{col}")
  49. indices = generate_indices(f, False)
  50. for i in indices:
  51. f.at[i] = 1
  52. expected = f.loc[i]
  53. tm.assert_almost_equal(expected, 1)
  54. class TestAtAndiAT:
  55. # at and iat tests that don't need Base class
  56. def test_float_index_at_iat(self):
  57. ser = Series([1, 2, 3], index=[0.1, 0.2, 0.3])
  58. for el, item in ser.items():
  59. assert ser.at[el] == item
  60. for i in range(len(ser)):
  61. assert ser.iat[i] == i + 1
  62. def test_at_iat_coercion(self):
  63. # as timestamp is not a tuple!
  64. dates = date_range("1/1/2000", periods=8)
  65. df = DataFrame(np.random.randn(8, 4), index=dates, columns=["A", "B", "C", "D"])
  66. s = df["A"]
  67. result = s.at[dates[5]]
  68. xp = s.values[5]
  69. assert result == xp
  70. @pytest.mark.parametrize(
  71. "ser, expected",
  72. [
  73. [
  74. Series(["2014-01-01", "2014-02-02"], dtype="datetime64[ns]"),
  75. Timestamp("2014-02-02"),
  76. ],
  77. [
  78. Series(["1 days", "2 days"], dtype="timedelta64[ns]"),
  79. Timedelta("2 days"),
  80. ],
  81. ],
  82. )
  83. def test_iloc_iat_coercion_datelike(self, indexer_ial, ser, expected):
  84. # GH 7729
  85. # make sure we are boxing the returns
  86. result = indexer_ial(ser)[1]
  87. assert result == expected
  88. def test_imethods_with_dups(self):
  89. # GH6493
  90. # iat/iloc with dups
  91. s = Series(range(5), index=[1, 1, 2, 2, 3], dtype="int64")
  92. result = s.iloc[2]
  93. assert result == 2
  94. result = s.iat[2]
  95. assert result == 2
  96. msg = "index 10 is out of bounds for axis 0 with size 5"
  97. with pytest.raises(IndexError, match=msg):
  98. s.iat[10]
  99. msg = "index -10 is out of bounds for axis 0 with size 5"
  100. with pytest.raises(IndexError, match=msg):
  101. s.iat[-10]
  102. result = s.iloc[[2, 3]]
  103. expected = Series([2, 3], [2, 2], dtype="int64")
  104. tm.assert_series_equal(result, expected)
  105. df = s.to_frame()
  106. result = df.iloc[2]
  107. expected = Series(2, index=[0], name=2)
  108. tm.assert_series_equal(result, expected)
  109. result = df.iat[2, 0]
  110. assert result == 2
  111. def test_frame_at_with_duplicate_axes(self):
  112. # GH#33041
  113. arr = np.random.randn(6).reshape(3, 2)
  114. df = DataFrame(arr, columns=["A", "A"])
  115. result = df.at[0, "A"]
  116. expected = df.iloc[0]
  117. tm.assert_series_equal(result, expected)
  118. result = df.T.at["A", 0]
  119. tm.assert_series_equal(result, expected)
  120. # setter
  121. df.at[1, "A"] = 2
  122. expected = Series([2.0, 2.0], index=["A", "A"], name=1)
  123. tm.assert_series_equal(df.iloc[1], expected)
  124. def test_at_getitem_dt64tz_values(self):
  125. # gh-15822
  126. df = DataFrame(
  127. {
  128. "name": ["John", "Anderson"],
  129. "date": [
  130. Timestamp(2017, 3, 13, 13, 32, 56),
  131. Timestamp(2017, 2, 16, 12, 10, 3),
  132. ],
  133. }
  134. )
  135. df["date"] = df["date"].dt.tz_localize("Asia/Shanghai")
  136. expected = Timestamp("2017-03-13 13:32:56+0800", tz="Asia/Shanghai")
  137. result = df.loc[0, "date"]
  138. assert result == expected
  139. result = df.at[0, "date"]
  140. assert result == expected
  141. def test_mixed_index_at_iat_loc_iloc_series(self):
  142. # GH 19860
  143. s = Series([1, 2, 3, 4, 5], index=["a", "b", "c", 1, 2])
  144. for el, item in s.items():
  145. assert s.at[el] == s.loc[el] == item
  146. for i in range(len(s)):
  147. assert s.iat[i] == s.iloc[i] == i + 1
  148. with pytest.raises(KeyError, match="^4$"):
  149. s.at[4]
  150. with pytest.raises(KeyError, match="^4$"):
  151. s.loc[4]
  152. def test_mixed_index_at_iat_loc_iloc_dataframe(self):
  153. # GH 19860
  154. df = DataFrame(
  155. [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], columns=["a", "b", "c", 1, 2]
  156. )
  157. for rowIdx, row in df.iterrows():
  158. for el, item in row.items():
  159. assert df.at[rowIdx, el] == df.loc[rowIdx, el] == item
  160. for row in range(2):
  161. for i in range(5):
  162. assert df.iat[row, i] == df.iloc[row, i] == row * 5 + i
  163. with pytest.raises(KeyError, match="^3$"):
  164. df.at[0, 3]
  165. with pytest.raises(KeyError, match="^3$"):
  166. df.loc[0, 3]
  167. def test_iat_setter_incompatible_assignment(self):
  168. # GH 23236
  169. result = DataFrame({"a": [0, 1], "b": [4, 5]})
  170. result.iat[0, 0] = None
  171. expected = DataFrame({"a": [None, 1], "b": [4, 5]})
  172. tm.assert_frame_equal(result, expected)
  173. def test_iat_dont_wrap_object_datetimelike():
  174. # GH#32809 .iat calls go through DataFrame._get_value, should not
  175. # call maybe_box_datetimelike
  176. dti = date_range("2016-01-01", periods=3)
  177. tdi = dti - dti
  178. ser = Series(dti.to_pydatetime(), dtype=object)
  179. ser2 = Series(tdi.to_pytimedelta(), dtype=object)
  180. df = DataFrame({"A": ser, "B": ser2})
  181. assert (df.dtypes == object).all()
  182. for result in [df.at[0, "A"], df.iat[0, 0], df.loc[0, "A"], df.iloc[0, 0]]:
  183. assert result is ser[0]
  184. assert isinstance(result, datetime)
  185. assert not isinstance(result, Timestamp)
  186. for result in [df.at[1, "B"], df.iat[1, 1], df.loc[1, "B"], df.iloc[1, 1]]:
  187. assert result is ser2[1]
  188. assert isinstance(result, timedelta)
  189. assert not isinstance(result, Timedelta)
  190. def test_at_with_tuple_index_get():
  191. # GH 26989
  192. # DataFrame.at getter works with Index of tuples
  193. df = DataFrame({"a": [1, 2]}, index=[(1, 2), (3, 4)])
  194. assert df.index.nlevels == 1
  195. assert df.at[(1, 2), "a"] == 1
  196. # Series.at getter works with Index of tuples
  197. series = df["a"]
  198. assert series.index.nlevels == 1
  199. assert series.at[(1, 2)] == 1
  200. def test_at_with_tuple_index_set():
  201. # GH 26989
  202. # DataFrame.at setter works with Index of tuples
  203. df = DataFrame({"a": [1, 2]}, index=[(1, 2), (3, 4)])
  204. assert df.index.nlevels == 1
  205. df.at[(1, 2), "a"] = 2
  206. assert df.at[(1, 2), "a"] == 2
  207. # Series.at setter works with Index of tuples
  208. series = df["a"]
  209. assert series.index.nlevels == 1
  210. series.at[1, 2] = 3
  211. assert series.at[1, 2] == 3
  212. class TestMultiIndexScalar:
  213. def test_multiindex_at_get(self):
  214. # GH 26989
  215. # DataFrame.at and DataFrame.loc getter works with MultiIndex
  216. df = DataFrame({"a": [1, 2]}, index=[[1, 2], [3, 4]])
  217. assert df.index.nlevels == 2
  218. assert df.at[(1, 3), "a"] == 1
  219. assert df.loc[(1, 3), "a"] == 1
  220. # Series.at and Series.loc getter works with MultiIndex
  221. series = df["a"]
  222. assert series.index.nlevels == 2
  223. assert series.at[1, 3] == 1
  224. assert series.loc[1, 3] == 1
  225. def test_multiindex_at_set(self):
  226. # GH 26989
  227. # DataFrame.at and DataFrame.loc setter works with MultiIndex
  228. df = DataFrame({"a": [1, 2]}, index=[[1, 2], [3, 4]])
  229. assert df.index.nlevels == 2
  230. df.at[(1, 3), "a"] = 3
  231. assert df.at[(1, 3), "a"] == 3
  232. df.loc[(1, 3), "a"] = 4
  233. assert df.loc[(1, 3), "a"] == 4
  234. # Series.at and Series.loc setter works with MultiIndex
  235. series = df["a"]
  236. assert series.index.nlevels == 2
  237. series.at[1, 3] = 5
  238. assert series.at[1, 3] == 5
  239. series.loc[1, 3] = 6
  240. assert series.loc[1, 3] == 6
  241. def test_multiindex_at_get_one_level(self):
  242. # GH#38053
  243. s2 = Series((0, 1), index=[[False, True]])
  244. result = s2.at[False]
  245. assert result == 0