test_textreader.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342
  1. """
  2. Tests the TextReader class in parsers.pyx, which
  3. is integral to the C engine in parsers.py
  4. """
  5. from io import (
  6. BytesIO,
  7. StringIO,
  8. )
  9. import numpy as np
  10. import pytest
  11. import pandas._libs.parsers as parser
  12. from pandas._libs.parsers import TextReader
  13. from pandas.errors import ParserWarning
  14. from pandas import DataFrame
  15. import pandas._testing as tm
  16. from pandas.io.parsers import (
  17. TextFileReader,
  18. read_csv,
  19. )
  20. from pandas.io.parsers.c_parser_wrapper import ensure_dtype_objs
  21. class TestTextReader:
  22. @pytest.fixture
  23. def csv_path(self, datapath):
  24. return datapath("io", "data", "csv", "test1.csv")
  25. def test_file_handle(self, csv_path):
  26. with open(csv_path, "rb") as f:
  27. reader = TextReader(f)
  28. reader.read()
  29. def test_file_handle_mmap(self, csv_path):
  30. # this was never using memory_map=True
  31. with open(csv_path, "rb") as f:
  32. reader = TextReader(f, header=None)
  33. reader.read()
  34. def test_StringIO(self, csv_path):
  35. with open(csv_path, "rb") as f:
  36. text = f.read()
  37. src = BytesIO(text)
  38. reader = TextReader(src, header=None)
  39. reader.read()
  40. def test_string_factorize(self):
  41. # should this be optional?
  42. data = "a\nb\na\nb\na"
  43. reader = TextReader(StringIO(data), header=None)
  44. result = reader.read()
  45. assert len(set(map(id, result[0]))) == 2
  46. def test_skipinitialspace(self):
  47. data = "a, b\na, b\na, b\na, b"
  48. reader = TextReader(StringIO(data), skipinitialspace=True, header=None)
  49. result = reader.read()
  50. tm.assert_numpy_array_equal(
  51. result[0], np.array(["a", "a", "a", "a"], dtype=np.object_)
  52. )
  53. tm.assert_numpy_array_equal(
  54. result[1], np.array(["b", "b", "b", "b"], dtype=np.object_)
  55. )
  56. def test_parse_booleans(self):
  57. data = "True\nFalse\nTrue\nTrue"
  58. reader = TextReader(StringIO(data), header=None)
  59. result = reader.read()
  60. assert result[0].dtype == np.bool_
  61. def test_delimit_whitespace(self):
  62. data = 'a b\na\t\t "b"\n"a"\t \t b'
  63. reader = TextReader(StringIO(data), delim_whitespace=True, header=None)
  64. result = reader.read()
  65. tm.assert_numpy_array_equal(
  66. result[0], np.array(["a", "a", "a"], dtype=np.object_)
  67. )
  68. tm.assert_numpy_array_equal(
  69. result[1], np.array(["b", "b", "b"], dtype=np.object_)
  70. )
  71. def test_embedded_newline(self):
  72. data = 'a\n"hello\nthere"\nthis'
  73. reader = TextReader(StringIO(data), header=None)
  74. result = reader.read()
  75. expected = np.array(["a", "hello\nthere", "this"], dtype=np.object_)
  76. tm.assert_numpy_array_equal(result[0], expected)
  77. def test_euro_decimal(self):
  78. data = "12345,67\n345,678"
  79. reader = TextReader(StringIO(data), delimiter=":", decimal=",", header=None)
  80. result = reader.read()
  81. expected = np.array([12345.67, 345.678])
  82. tm.assert_almost_equal(result[0], expected)
  83. def test_integer_thousands(self):
  84. data = "123,456\n12,500"
  85. reader = TextReader(StringIO(data), delimiter=":", thousands=",", header=None)
  86. result = reader.read()
  87. expected = np.array([123456, 12500], dtype=np.int64)
  88. tm.assert_almost_equal(result[0], expected)
  89. def test_integer_thousands_alt(self):
  90. data = "123.456\n12.500"
  91. reader = TextFileReader(
  92. StringIO(data), delimiter=":", thousands=".", header=None
  93. )
  94. result = reader.read()
  95. expected = DataFrame([123456, 12500])
  96. tm.assert_frame_equal(result, expected)
  97. def test_skip_bad_lines(self):
  98. # too many lines, see #2430 for why
  99. data = "a:b:c\nd:e:f\ng:h:i\nj:k:l:m\nl:m:n\no:p:q:r"
  100. reader = TextReader(StringIO(data), delimiter=":", header=None)
  101. msg = r"Error tokenizing data\. C error: Expected 3 fields in line 4, saw 4"
  102. with pytest.raises(parser.ParserError, match=msg):
  103. reader.read()
  104. reader = TextReader(
  105. StringIO(data), delimiter=":", header=None, on_bad_lines=2 # Skip
  106. )
  107. result = reader.read()
  108. expected = {
  109. 0: np.array(["a", "d", "g", "l"], dtype=object),
  110. 1: np.array(["b", "e", "h", "m"], dtype=object),
  111. 2: np.array(["c", "f", "i", "n"], dtype=object),
  112. }
  113. assert_array_dicts_equal(result, expected)
  114. with tm.assert_produces_warning(ParserWarning, match="Skipping line"):
  115. reader = TextReader(
  116. StringIO(data), delimiter=":", header=None, on_bad_lines=1 # Warn
  117. )
  118. reader.read()
  119. def test_header_not_enough_lines(self):
  120. data = "skip this\nskip this\na,b,c\n1,2,3\n4,5,6"
  121. reader = TextReader(StringIO(data), delimiter=",", header=2)
  122. header = reader.header
  123. expected = [["a", "b", "c"]]
  124. assert header == expected
  125. recs = reader.read()
  126. expected = {
  127. 0: np.array([1, 4], dtype=np.int64),
  128. 1: np.array([2, 5], dtype=np.int64),
  129. 2: np.array([3, 6], dtype=np.int64),
  130. }
  131. assert_array_dicts_equal(recs, expected)
  132. def test_escapechar(self):
  133. data = '\\"hello world"\n\\"hello world"\n\\"hello world"'
  134. reader = TextReader(StringIO(data), delimiter=",", header=None, escapechar="\\")
  135. result = reader.read()
  136. expected = {0: np.array(['"hello world"'] * 3, dtype=object)}
  137. assert_array_dicts_equal(result, expected)
  138. def test_eof_has_eol(self):
  139. # handling of new line at EOF
  140. pass
  141. def test_na_substitution(self):
  142. pass
  143. def test_numpy_string_dtype(self):
  144. data = """\
  145. a,1
  146. aa,2
  147. aaa,3
  148. aaaa,4
  149. aaaaa,5"""
  150. def _make_reader(**kwds):
  151. if "dtype" in kwds:
  152. kwds["dtype"] = ensure_dtype_objs(kwds["dtype"])
  153. return TextReader(StringIO(data), delimiter=",", header=None, **kwds)
  154. reader = _make_reader(dtype="S5,i4")
  155. result = reader.read()
  156. assert result[0].dtype == "S5"
  157. ex_values = np.array(["a", "aa", "aaa", "aaaa", "aaaaa"], dtype="S5")
  158. assert (result[0] == ex_values).all()
  159. assert result[1].dtype == "i4"
  160. reader = _make_reader(dtype="S4")
  161. result = reader.read()
  162. assert result[0].dtype == "S4"
  163. ex_values = np.array(["a", "aa", "aaa", "aaaa", "aaaa"], dtype="S4")
  164. assert (result[0] == ex_values).all()
  165. assert result[1].dtype == "S4"
  166. def test_pass_dtype(self):
  167. data = """\
  168. one,two
  169. 1,a
  170. 2,b
  171. 3,c
  172. 4,d"""
  173. def _make_reader(**kwds):
  174. if "dtype" in kwds:
  175. kwds["dtype"] = ensure_dtype_objs(kwds["dtype"])
  176. return TextReader(StringIO(data), delimiter=",", **kwds)
  177. reader = _make_reader(dtype={"one": "u1", 1: "S1"})
  178. result = reader.read()
  179. assert result[0].dtype == "u1"
  180. assert result[1].dtype == "S1"
  181. reader = _make_reader(dtype={"one": np.uint8, 1: object})
  182. result = reader.read()
  183. assert result[0].dtype == "u1"
  184. assert result[1].dtype == "O"
  185. reader = _make_reader(dtype={"one": np.dtype("u1"), 1: np.dtype("O")})
  186. result = reader.read()
  187. assert result[0].dtype == "u1"
  188. assert result[1].dtype == "O"
  189. def test_usecols(self):
  190. data = """\
  191. a,b,c
  192. 1,2,3
  193. 4,5,6
  194. 7,8,9
  195. 10,11,12"""
  196. def _make_reader(**kwds):
  197. return TextReader(StringIO(data), delimiter=",", **kwds)
  198. reader = _make_reader(usecols=(1, 2))
  199. result = reader.read()
  200. exp = _make_reader().read()
  201. assert len(result) == 2
  202. assert (result[1] == exp[1]).all()
  203. assert (result[2] == exp[2]).all()
  204. @pytest.mark.parametrize(
  205. "text, kwargs",
  206. [
  207. ("a,b,c\r1,2,3\r4,5,6\r7,8,9\r10,11,12", {"delimiter": ","}),
  208. (
  209. "a b c\r1 2 3\r4 5 6\r7 8 9\r10 11 12",
  210. {"delim_whitespace": True},
  211. ),
  212. ("a,b,c\r1,2,3\r4,5,6\r,88,9\r10,11,12", {"delimiter": ","}),
  213. (
  214. (
  215. "A,B,C,D,E,F,G,H,I,J,K,L,M,N,O\r"
  216. "AAAAA,BBBBB,0,0,0,0,0,0,0,0,0,0,0,0,0\r"
  217. ",BBBBB,0,0,0,0,0,0,0,0,0,0,0,0,0"
  218. ),
  219. {"delimiter": ","},
  220. ),
  221. ("A B C\r 2 3\r4 5 6", {"delim_whitespace": True}),
  222. ("A B C\r2 3\r4 5 6", {"delim_whitespace": True}),
  223. ],
  224. )
  225. def test_cr_delimited(self, text, kwargs):
  226. nice_text = text.replace("\r", "\r\n")
  227. result = TextReader(StringIO(text), **kwargs).read()
  228. expected = TextReader(StringIO(nice_text), **kwargs).read()
  229. assert_array_dicts_equal(result, expected)
  230. def test_empty_field_eof(self):
  231. data = "a,b,c\n1,2,3\n4,,"
  232. result = TextReader(StringIO(data), delimiter=",").read()
  233. expected = {
  234. 0: np.array([1, 4], dtype=np.int64),
  235. 1: np.array(["2", ""], dtype=object),
  236. 2: np.array(["3", ""], dtype=object),
  237. }
  238. assert_array_dicts_equal(result, expected)
  239. @pytest.mark.parametrize("repeat", range(10))
  240. def test_empty_field_eof_mem_access_bug(self, repeat):
  241. # GH5664
  242. a = DataFrame([["b"], [np.nan]], columns=["a"], index=["a", "c"])
  243. b = DataFrame([[1, 1, 1, 0], [1, 1, 1, 0]], columns=list("abcd"), index=[1, 1])
  244. c = DataFrame(
  245. [
  246. [1, 2, 3, 4],
  247. [6, np.nan, np.nan, np.nan],
  248. [8, 9, 10, 11],
  249. [13, 14, np.nan, np.nan],
  250. ],
  251. columns=list("abcd"),
  252. index=[0, 5, 7, 12],
  253. )
  254. df = read_csv(StringIO("a,b\nc\n"), skiprows=0, names=["a"], engine="c")
  255. tm.assert_frame_equal(df, a)
  256. df = read_csv(
  257. StringIO("1,1,1,1,0\n" * 2 + "\n" * 2), names=list("abcd"), engine="c"
  258. )
  259. tm.assert_frame_equal(df, b)
  260. df = read_csv(
  261. StringIO("0,1,2,3,4\n5,6\n7,8,9,10,11\n12,13,14"),
  262. names=list("abcd"),
  263. engine="c",
  264. )
  265. tm.assert_frame_equal(df, c)
  266. def test_empty_csv_input(self):
  267. # GH14867
  268. with read_csv(
  269. StringIO(), chunksize=20, header=None, names=["a", "b", "c"]
  270. ) as df:
  271. assert isinstance(df, TextFileReader)
  272. def assert_array_dicts_equal(left, right):
  273. for k, v in left.items():
  274. tm.assert_numpy_array_equal(np.asarray(v), np.asarray(right[k]))