common.py 40 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267
  1. """Common IO api utilities"""
  2. from __future__ import annotations
  3. from abc import (
  4. ABC,
  5. abstractmethod,
  6. )
  7. import codecs
  8. from collections import defaultdict
  9. from collections.abc import (
  10. Hashable,
  11. Mapping,
  12. Sequence,
  13. )
  14. import dataclasses
  15. import functools
  16. import gzip
  17. from io import (
  18. BufferedIOBase,
  19. BytesIO,
  20. RawIOBase,
  21. StringIO,
  22. TextIOBase,
  23. TextIOWrapper,
  24. )
  25. import mmap
  26. import os
  27. from pathlib import Path
  28. import re
  29. import tarfile
  30. from typing import (
  31. IO,
  32. TYPE_CHECKING,
  33. Any,
  34. AnyStr,
  35. DefaultDict,
  36. Generic,
  37. Literal,
  38. TypeVar,
  39. cast,
  40. overload,
  41. )
  42. from urllib.parse import (
  43. urljoin,
  44. urlparse as parse_url,
  45. uses_netloc,
  46. uses_params,
  47. uses_relative,
  48. )
  49. import warnings
  50. import zipfile
  51. from pandas._typing import (
  52. BaseBuffer,
  53. ReadCsvBuffer,
  54. )
  55. from pandas.compat import (
  56. get_bz2_file,
  57. get_lzma_file,
  58. )
  59. from pandas.compat._optional import import_optional_dependency
  60. from pandas.util._decorators import doc
  61. from pandas.util._exceptions import find_stack_level
  62. from pandas.core.dtypes.common import (
  63. is_bool,
  64. is_file_like,
  65. is_integer,
  66. is_list_like,
  67. )
  68. from pandas.core.dtypes.generic import ABCMultiIndex
  69. from pandas.core.shared_docs import _shared_docs
  70. _VALID_URLS = set(uses_relative + uses_netloc + uses_params)
  71. _VALID_URLS.discard("")
  72. _RFC_3986_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9+\-+.]*://")
  73. BaseBufferT = TypeVar("BaseBufferT", bound=BaseBuffer)
  74. if TYPE_CHECKING:
  75. from types import TracebackType
  76. from pandas._typing import (
  77. CompressionDict,
  78. CompressionOptions,
  79. FilePath,
  80. ReadBuffer,
  81. StorageOptions,
  82. WriteBuffer,
  83. )
  84. from pandas import MultiIndex
  85. @dataclasses.dataclass
  86. class IOArgs:
  87. """
  88. Return value of io/common.py:_get_filepath_or_buffer.
  89. """
  90. filepath_or_buffer: str | BaseBuffer
  91. encoding: str
  92. mode: str
  93. compression: CompressionDict
  94. should_close: bool = False
  95. @dataclasses.dataclass
  96. class IOHandles(Generic[AnyStr]):
  97. """
  98. Return value of io/common.py:get_handle
  99. Can be used as a context manager.
  100. This is used to easily close created buffers and to handle corner cases when
  101. TextIOWrapper is inserted.
  102. handle: The file handle to be used.
  103. created_handles: All file handles that are created by get_handle
  104. is_wrapped: Whether a TextIOWrapper needs to be detached.
  105. """
  106. # handle might not implement the IO-interface
  107. handle: IO[AnyStr]
  108. compression: CompressionDict
  109. created_handles: list[IO[bytes] | IO[str]] = dataclasses.field(default_factory=list)
  110. is_wrapped: bool = False
  111. def close(self) -> None:
  112. """
  113. Close all created buffers.
  114. Note: If a TextIOWrapper was inserted, it is flushed and detached to
  115. avoid closing the potentially user-created buffer.
  116. """
  117. if self.is_wrapped:
  118. assert isinstance(self.handle, TextIOWrapper)
  119. self.handle.flush()
  120. self.handle.detach()
  121. self.created_handles.remove(self.handle)
  122. for handle in self.created_handles:
  123. handle.close()
  124. self.created_handles = []
  125. self.is_wrapped = False
  126. def __enter__(self) -> IOHandles[AnyStr]:
  127. return self
  128. def __exit__(
  129. self,
  130. exc_type: type[BaseException] | None,
  131. exc_value: BaseException | None,
  132. traceback: TracebackType | None,
  133. ) -> None:
  134. self.close()
  135. def is_url(url: object) -> bool:
  136. """
  137. Check to see if a URL has a valid protocol.
  138. Parameters
  139. ----------
  140. url : str or unicode
  141. Returns
  142. -------
  143. isurl : bool
  144. If `url` has a valid protocol return True otherwise False.
  145. """
  146. if not isinstance(url, str):
  147. return False
  148. return parse_url(url).scheme in _VALID_URLS
  149. @overload
  150. def _expand_user(filepath_or_buffer: str) -> str:
  151. ...
  152. @overload
  153. def _expand_user(filepath_or_buffer: BaseBufferT) -> BaseBufferT:
  154. ...
  155. def _expand_user(filepath_or_buffer: str | BaseBufferT) -> str | BaseBufferT:
  156. """
  157. Return the argument with an initial component of ~ or ~user
  158. replaced by that user's home directory.
  159. Parameters
  160. ----------
  161. filepath_or_buffer : object to be converted if possible
  162. Returns
  163. -------
  164. expanded_filepath_or_buffer : an expanded filepath or the
  165. input if not expandable
  166. """
  167. if isinstance(filepath_or_buffer, str):
  168. return os.path.expanduser(filepath_or_buffer)
  169. return filepath_or_buffer
  170. def validate_header_arg(header: object) -> None:
  171. if header is None:
  172. return
  173. if is_integer(header):
  174. header = cast(int, header)
  175. if header < 0:
  176. # GH 27779
  177. raise ValueError(
  178. "Passing negative integer to header is invalid. "
  179. "For no header, use header=None instead"
  180. )
  181. return
  182. if is_list_like(header, allow_sets=False):
  183. header = cast(Sequence, header)
  184. if not all(map(is_integer, header)):
  185. raise ValueError("header must be integer or list of integers")
  186. if any(i < 0 for i in header):
  187. raise ValueError("cannot specify multi-index header with negative integers")
  188. return
  189. if is_bool(header):
  190. raise TypeError(
  191. "Passing a bool to header is invalid. Use header=None for no header or "
  192. "header=int or list-like of ints to specify "
  193. "the row(s) making up the column names"
  194. )
  195. # GH 16338
  196. raise ValueError("header must be integer or list of integers")
  197. @overload
  198. def stringify_path(filepath_or_buffer: FilePath, convert_file_like: bool = ...) -> str:
  199. ...
  200. @overload
  201. def stringify_path(
  202. filepath_or_buffer: BaseBufferT, convert_file_like: bool = ...
  203. ) -> BaseBufferT:
  204. ...
  205. def stringify_path(
  206. filepath_or_buffer: FilePath | BaseBufferT,
  207. convert_file_like: bool = False,
  208. ) -> str | BaseBufferT:
  209. """
  210. Attempt to convert a path-like object to a string.
  211. Parameters
  212. ----------
  213. filepath_or_buffer : object to be converted
  214. Returns
  215. -------
  216. str_filepath_or_buffer : maybe a string version of the object
  217. Notes
  218. -----
  219. Objects supporting the fspath protocol are coerced
  220. according to its __fspath__ method.
  221. Any other object is passed through unchanged, which includes bytes,
  222. strings, buffers, or anything else that's not even path-like.
  223. """
  224. if not convert_file_like and is_file_like(filepath_or_buffer):
  225. # GH 38125: some fsspec objects implement os.PathLike but have already opened a
  226. # file. This prevents opening the file a second time. infer_compression calls
  227. # this function with convert_file_like=True to infer the compression.
  228. return cast(BaseBufferT, filepath_or_buffer)
  229. if isinstance(filepath_or_buffer, os.PathLike):
  230. filepath_or_buffer = filepath_or_buffer.__fspath__()
  231. return _expand_user(filepath_or_buffer)
  232. def urlopen(*args, **kwargs):
  233. """
  234. Lazy-import wrapper for stdlib urlopen, as that imports a big chunk of
  235. the stdlib.
  236. """
  237. import urllib.request
  238. return urllib.request.urlopen(*args, **kwargs)
  239. def is_fsspec_url(url: FilePath | BaseBuffer) -> bool:
  240. """
  241. Returns true if the given URL looks like
  242. something fsspec can handle
  243. """
  244. return (
  245. isinstance(url, str)
  246. and bool(_RFC_3986_PATTERN.match(url))
  247. and not url.startswith(("http://", "https://"))
  248. )
  249. @doc(
  250. storage_options=_shared_docs["storage_options"],
  251. compression_options=_shared_docs["compression_options"] % "filepath_or_buffer",
  252. )
  253. def _get_filepath_or_buffer(
  254. filepath_or_buffer: FilePath | BaseBuffer,
  255. encoding: str = "utf-8",
  256. compression: CompressionOptions | None = None,
  257. mode: str = "r",
  258. storage_options: StorageOptions | None = None,
  259. ) -> IOArgs:
  260. """
  261. If the filepath_or_buffer is a url, translate and return the buffer.
  262. Otherwise passthrough.
  263. Parameters
  264. ----------
  265. filepath_or_buffer : a url, filepath (str, py.path.local or pathlib.Path),
  266. or buffer
  267. {compression_options}
  268. .. versionchanged:: 1.4.0 Zstandard support.
  269. encoding : the encoding to use to decode bytes, default is 'utf-8'
  270. mode : str, optional
  271. {storage_options}
  272. Returns the dataclass IOArgs.
  273. """
  274. filepath_or_buffer = stringify_path(filepath_or_buffer)
  275. # handle compression dict
  276. compression_method, compression = get_compression_method(compression)
  277. compression_method = infer_compression(filepath_or_buffer, compression_method)
  278. # GH21227 internal compression is not used for non-binary handles.
  279. if compression_method and hasattr(filepath_or_buffer, "write") and "b" not in mode:
  280. warnings.warn(
  281. "compression has no effect when passing a non-binary object as input.",
  282. RuntimeWarning,
  283. stacklevel=find_stack_level(),
  284. )
  285. compression_method = None
  286. compression = dict(compression, method=compression_method)
  287. # bz2 and xz do not write the byte order mark for utf-16 and utf-32
  288. # print a warning when writing such files
  289. if (
  290. "w" in mode
  291. and compression_method in ["bz2", "xz"]
  292. and encoding in ["utf-16", "utf-32"]
  293. ):
  294. warnings.warn(
  295. f"{compression} will not write the byte order mark for {encoding}",
  296. UnicodeWarning,
  297. stacklevel=find_stack_level(),
  298. )
  299. # Use binary mode when converting path-like objects to file-like objects (fsspec)
  300. # except when text mode is explicitly requested. The original mode is returned if
  301. # fsspec is not used.
  302. fsspec_mode = mode
  303. if "t" not in fsspec_mode and "b" not in fsspec_mode:
  304. fsspec_mode += "b"
  305. if isinstance(filepath_or_buffer, str) and is_url(filepath_or_buffer):
  306. # TODO: fsspec can also handle HTTP via requests, but leaving this
  307. # unchanged. using fsspec appears to break the ability to infer if the
  308. # server responded with gzipped data
  309. storage_options = storage_options or {}
  310. # waiting until now for importing to match intended lazy logic of
  311. # urlopen function defined elsewhere in this module
  312. import urllib.request
  313. # assuming storage_options is to be interpreted as headers
  314. req_info = urllib.request.Request(filepath_or_buffer, headers=storage_options)
  315. with urlopen(req_info) as req:
  316. content_encoding = req.headers.get("Content-Encoding", None)
  317. if content_encoding == "gzip":
  318. # Override compression based on Content-Encoding header
  319. compression = {"method": "gzip"}
  320. reader = BytesIO(req.read())
  321. return IOArgs(
  322. filepath_or_buffer=reader,
  323. encoding=encoding,
  324. compression=compression,
  325. should_close=True,
  326. mode=fsspec_mode,
  327. )
  328. if is_fsspec_url(filepath_or_buffer):
  329. assert isinstance(
  330. filepath_or_buffer, str
  331. ) # just to appease mypy for this branch
  332. # two special-case s3-like protocols; these have special meaning in Hadoop,
  333. # but are equivalent to just "s3" from fsspec's point of view
  334. # cc #11071
  335. if filepath_or_buffer.startswith("s3a://"):
  336. filepath_or_buffer = filepath_or_buffer.replace("s3a://", "s3://")
  337. if filepath_or_buffer.startswith("s3n://"):
  338. filepath_or_buffer = filepath_or_buffer.replace("s3n://", "s3://")
  339. fsspec = import_optional_dependency("fsspec")
  340. # If botocore is installed we fallback to reading with anon=True
  341. # to allow reads from public buckets
  342. err_types_to_retry_with_anon: list[Any] = []
  343. try:
  344. import_optional_dependency("botocore")
  345. from botocore.exceptions import (
  346. ClientError,
  347. NoCredentialsError,
  348. )
  349. err_types_to_retry_with_anon = [
  350. ClientError,
  351. NoCredentialsError,
  352. PermissionError,
  353. ]
  354. except ImportError:
  355. pass
  356. try:
  357. file_obj = fsspec.open(
  358. filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
  359. ).open()
  360. # GH 34626 Reads from Public Buckets without Credentials needs anon=True
  361. except tuple(err_types_to_retry_with_anon):
  362. if storage_options is None:
  363. storage_options = {"anon": True}
  364. else:
  365. # don't mutate user input.
  366. storage_options = dict(storage_options)
  367. storage_options["anon"] = True
  368. file_obj = fsspec.open(
  369. filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
  370. ).open()
  371. return IOArgs(
  372. filepath_or_buffer=file_obj,
  373. encoding=encoding,
  374. compression=compression,
  375. should_close=True,
  376. mode=fsspec_mode,
  377. )
  378. elif storage_options:
  379. raise ValueError(
  380. "storage_options passed with file object or non-fsspec file path"
  381. )
  382. if isinstance(filepath_or_buffer, (str, bytes, mmap.mmap)):
  383. return IOArgs(
  384. filepath_or_buffer=_expand_user(filepath_or_buffer),
  385. encoding=encoding,
  386. compression=compression,
  387. should_close=False,
  388. mode=mode,
  389. )
  390. # is_file_like requires (read | write) & __iter__ but __iter__ is only
  391. # needed for read_csv(engine=python)
  392. if not (
  393. hasattr(filepath_or_buffer, "read") or hasattr(filepath_or_buffer, "write")
  394. ):
  395. msg = f"Invalid file path or buffer object type: {type(filepath_or_buffer)}"
  396. raise ValueError(msg)
  397. return IOArgs(
  398. filepath_or_buffer=filepath_or_buffer,
  399. encoding=encoding,
  400. compression=compression,
  401. should_close=False,
  402. mode=mode,
  403. )
  404. def file_path_to_url(path: str) -> str:
  405. """
  406. converts an absolute native path to a FILE URL.
  407. Parameters
  408. ----------
  409. path : a path in native format
  410. Returns
  411. -------
  412. a valid FILE URL
  413. """
  414. # lazify expensive import (~30ms)
  415. from urllib.request import pathname2url
  416. return urljoin("file:", pathname2url(path))
  417. extension_to_compression = {
  418. ".tar": "tar",
  419. ".tar.gz": "tar",
  420. ".tar.bz2": "tar",
  421. ".tar.xz": "tar",
  422. ".gz": "gzip",
  423. ".bz2": "bz2",
  424. ".zip": "zip",
  425. ".xz": "xz",
  426. ".zst": "zstd",
  427. }
  428. _supported_compressions = set(extension_to_compression.values())
  429. def get_compression_method(
  430. compression: CompressionOptions,
  431. ) -> tuple[str | None, CompressionDict]:
  432. """
  433. Simplifies a compression argument to a compression method string and
  434. a mapping containing additional arguments.
  435. Parameters
  436. ----------
  437. compression : str or mapping
  438. If string, specifies the compression method. If mapping, value at key
  439. 'method' specifies compression method.
  440. Returns
  441. -------
  442. tuple of ({compression method}, Optional[str]
  443. {compression arguments}, Dict[str, Any])
  444. Raises
  445. ------
  446. ValueError on mapping missing 'method' key
  447. """
  448. compression_method: str | None
  449. if isinstance(compression, Mapping):
  450. compression_args = dict(compression)
  451. try:
  452. compression_method = compression_args.pop("method")
  453. except KeyError as err:
  454. raise ValueError("If mapping, compression must have key 'method'") from err
  455. else:
  456. compression_args = {}
  457. compression_method = compression
  458. return compression_method, compression_args
  459. @doc(compression_options=_shared_docs["compression_options"] % "filepath_or_buffer")
  460. def infer_compression(
  461. filepath_or_buffer: FilePath | BaseBuffer, compression: str | None
  462. ) -> str | None:
  463. """
  464. Get the compression method for filepath_or_buffer. If compression='infer',
  465. the inferred compression method is returned. Otherwise, the input
  466. compression method is returned unchanged, unless it's invalid, in which
  467. case an error is raised.
  468. Parameters
  469. ----------
  470. filepath_or_buffer : str or file handle
  471. File path or object.
  472. {compression_options}
  473. .. versionchanged:: 1.4.0 Zstandard support.
  474. Returns
  475. -------
  476. string or None
  477. Raises
  478. ------
  479. ValueError on invalid compression specified.
  480. """
  481. if compression is None:
  482. return None
  483. # Infer compression
  484. if compression == "infer":
  485. # Convert all path types (e.g. pathlib.Path) to strings
  486. filepath_or_buffer = stringify_path(filepath_or_buffer, convert_file_like=True)
  487. if not isinstance(filepath_or_buffer, str):
  488. # Cannot infer compression of a buffer, assume no compression
  489. return None
  490. # Infer compression from the filename/URL extension
  491. for extension, compression in extension_to_compression.items():
  492. if filepath_or_buffer.lower().endswith(extension):
  493. return compression
  494. return None
  495. # Compression has been specified. Check that it's valid
  496. if compression in _supported_compressions:
  497. return compression
  498. valid = ["infer", None] + sorted(_supported_compressions)
  499. msg = (
  500. f"Unrecognized compression type: {compression}\n"
  501. f"Valid compression types are {valid}"
  502. )
  503. raise ValueError(msg)
  504. def check_parent_directory(path: Path | str) -> None:
  505. """
  506. Check if parent directory of a file exists, raise OSError if it does not
  507. Parameters
  508. ----------
  509. path: Path or str
  510. Path to check parent directory of
  511. """
  512. parent = Path(path).parent
  513. if not parent.is_dir():
  514. raise OSError(rf"Cannot save file into a non-existent directory: '{parent}'")
  515. @overload
  516. def get_handle(
  517. path_or_buf: FilePath | BaseBuffer,
  518. mode: str,
  519. *,
  520. encoding: str | None = ...,
  521. compression: CompressionOptions = ...,
  522. memory_map: bool = ...,
  523. is_text: Literal[False],
  524. errors: str | None = ...,
  525. storage_options: StorageOptions = ...,
  526. ) -> IOHandles[bytes]:
  527. ...
  528. @overload
  529. def get_handle(
  530. path_or_buf: FilePath | BaseBuffer,
  531. mode: str,
  532. *,
  533. encoding: str | None = ...,
  534. compression: CompressionOptions = ...,
  535. memory_map: bool = ...,
  536. is_text: Literal[True] = ...,
  537. errors: str | None = ...,
  538. storage_options: StorageOptions = ...,
  539. ) -> IOHandles[str]:
  540. ...
  541. @overload
  542. def get_handle(
  543. path_or_buf: FilePath | BaseBuffer,
  544. mode: str,
  545. *,
  546. encoding: str | None = ...,
  547. compression: CompressionOptions = ...,
  548. memory_map: bool = ...,
  549. is_text: bool = ...,
  550. errors: str | None = ...,
  551. storage_options: StorageOptions = ...,
  552. ) -> IOHandles[str] | IOHandles[bytes]:
  553. ...
  554. @doc(compression_options=_shared_docs["compression_options"] % "path_or_buf")
  555. def get_handle(
  556. path_or_buf: FilePath | BaseBuffer,
  557. mode: str,
  558. *,
  559. encoding: str | None = None,
  560. compression: CompressionOptions | None = None,
  561. memory_map: bool = False,
  562. is_text: bool = True,
  563. errors: str | None = None,
  564. storage_options: StorageOptions | None = None,
  565. ) -> IOHandles[str] | IOHandles[bytes]:
  566. """
  567. Get file handle for given path/buffer and mode.
  568. Parameters
  569. ----------
  570. path_or_buf : str or file handle
  571. File path or object.
  572. mode : str
  573. Mode to open path_or_buf with.
  574. encoding : str or None
  575. Encoding to use.
  576. {compression_options}
  577. May be a dict with key 'method' as compression mode
  578. and other keys as compression options if compression
  579. mode is 'zip'.
  580. Passing compression options as keys in dict is
  581. supported for compression modes 'gzip', 'bz2', 'zstd' and 'zip'.
  582. .. versionchanged:: 1.4.0 Zstandard support.
  583. memory_map : bool, default False
  584. See parsers._parser_params for more information. Only used by read_csv.
  585. is_text : bool, default True
  586. Whether the type of the content passed to the file/buffer is string or
  587. bytes. This is not the same as `"b" not in mode`. If a string content is
  588. passed to a binary file/buffer, a wrapper is inserted.
  589. errors : str, default 'strict'
  590. Specifies how encoding and decoding errors are to be handled.
  591. See the errors argument for :func:`open` for a full list
  592. of options.
  593. storage_options: StorageOptions = None
  594. Passed to _get_filepath_or_buffer
  595. Returns the dataclass IOHandles
  596. """
  597. # Windows does not default to utf-8. Set to utf-8 for a consistent behavior
  598. encoding = encoding or "utf-8"
  599. errors = errors or "strict"
  600. # read_csv does not know whether the buffer is opened in binary/text mode
  601. if _is_binary_mode(path_or_buf, mode) and "b" not in mode:
  602. mode += "b"
  603. # validate encoding and errors
  604. codecs.lookup(encoding)
  605. if isinstance(errors, str):
  606. codecs.lookup_error(errors)
  607. # open URLs
  608. ioargs = _get_filepath_or_buffer(
  609. path_or_buf,
  610. encoding=encoding,
  611. compression=compression,
  612. mode=mode,
  613. storage_options=storage_options,
  614. )
  615. handle = ioargs.filepath_or_buffer
  616. handles: list[BaseBuffer]
  617. # memory mapping needs to be the first step
  618. # only used for read_csv
  619. handle, memory_map, handles = _maybe_memory_map(handle, memory_map)
  620. is_path = isinstance(handle, str)
  621. compression_args = dict(ioargs.compression)
  622. compression = compression_args.pop("method")
  623. # Only for write methods
  624. if "r" not in mode and is_path:
  625. check_parent_directory(str(handle))
  626. if compression:
  627. if compression != "zstd":
  628. # compression libraries do not like an explicit text-mode
  629. ioargs.mode = ioargs.mode.replace("t", "")
  630. elif compression == "zstd" and "b" not in ioargs.mode:
  631. # python-zstandard defaults to text mode, but we always expect
  632. # compression libraries to use binary mode.
  633. ioargs.mode += "b"
  634. # GZ Compression
  635. if compression == "gzip":
  636. if isinstance(handle, str):
  637. # error: Incompatible types in assignment (expression has type
  638. # "GzipFile", variable has type "Union[str, BaseBuffer]")
  639. handle = gzip.GzipFile( # type: ignore[assignment]
  640. filename=handle,
  641. mode=ioargs.mode,
  642. **compression_args,
  643. )
  644. else:
  645. handle = gzip.GzipFile(
  646. # No overload variant of "GzipFile" matches argument types
  647. # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
  648. fileobj=handle, # type: ignore[call-overload]
  649. mode=ioargs.mode,
  650. **compression_args,
  651. )
  652. # BZ Compression
  653. elif compression == "bz2":
  654. # Overload of "BZ2File" to handle pickle protocol 5
  655. # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
  656. handle = get_bz2_file()( # type: ignore[call-overload]
  657. handle,
  658. mode=ioargs.mode,
  659. **compression_args,
  660. )
  661. # ZIP Compression
  662. elif compression == "zip":
  663. # error: Argument 1 to "_BytesZipFile" has incompatible type
  664. # "Union[str, BaseBuffer]"; expected "Union[Union[str, PathLike[str]],
  665. # ReadBuffer[bytes], WriteBuffer[bytes]]"
  666. handle = _BytesZipFile(
  667. handle, ioargs.mode, **compression_args # type: ignore[arg-type]
  668. )
  669. if handle.buffer.mode == "r":
  670. handles.append(handle)
  671. zip_names = handle.buffer.namelist()
  672. if len(zip_names) == 1:
  673. handle = handle.buffer.open(zip_names.pop())
  674. elif not zip_names:
  675. raise ValueError(f"Zero files found in ZIP file {path_or_buf}")
  676. else:
  677. raise ValueError(
  678. "Multiple files found in ZIP file. "
  679. f"Only one file per ZIP: {zip_names}"
  680. )
  681. # TAR Encoding
  682. elif compression == "tar":
  683. compression_args.setdefault("mode", ioargs.mode)
  684. if isinstance(handle, str):
  685. handle = _BytesTarFile(name=handle, **compression_args)
  686. else:
  687. # error: Argument "fileobj" to "_BytesTarFile" has incompatible
  688. # type "BaseBuffer"; expected "Union[ReadBuffer[bytes],
  689. # WriteBuffer[bytes], None]"
  690. handle = _BytesTarFile(
  691. fileobj=handle, **compression_args # type: ignore[arg-type]
  692. )
  693. assert isinstance(handle, _BytesTarFile)
  694. if "r" in handle.buffer.mode:
  695. handles.append(handle)
  696. files = handle.buffer.getnames()
  697. if len(files) == 1:
  698. file = handle.buffer.extractfile(files[0])
  699. assert file is not None
  700. handle = file
  701. elif not files:
  702. raise ValueError(f"Zero files found in TAR archive {path_or_buf}")
  703. else:
  704. raise ValueError(
  705. "Multiple files found in TAR archive. "
  706. f"Only one file per TAR archive: {files}"
  707. )
  708. # XZ Compression
  709. elif compression == "xz":
  710. # error: Argument 1 to "LZMAFile" has incompatible type "Union[str,
  711. # BaseBuffer]"; expected "Optional[Union[Union[str, bytes, PathLike[str],
  712. # PathLike[bytes]], IO[bytes]], None]"
  713. handle = get_lzma_file()(
  714. handle, ioargs.mode, **compression_args # type: ignore[arg-type]
  715. )
  716. # Zstd Compression
  717. elif compression == "zstd":
  718. zstd = import_optional_dependency("zstandard")
  719. if "r" in ioargs.mode:
  720. open_args = {"dctx": zstd.ZstdDecompressor(**compression_args)}
  721. else:
  722. open_args = {"cctx": zstd.ZstdCompressor(**compression_args)}
  723. handle = zstd.open(
  724. handle,
  725. mode=ioargs.mode,
  726. **open_args,
  727. )
  728. # Unrecognized Compression
  729. else:
  730. msg = f"Unrecognized compression type: {compression}"
  731. raise ValueError(msg)
  732. assert not isinstance(handle, str)
  733. handles.append(handle)
  734. elif isinstance(handle, str):
  735. # Check whether the filename is to be opened in binary mode.
  736. # Binary mode does not support 'encoding' and 'newline'.
  737. if ioargs.encoding and "b" not in ioargs.mode:
  738. # Encoding
  739. handle = open(
  740. handle,
  741. ioargs.mode,
  742. encoding=ioargs.encoding,
  743. errors=errors,
  744. newline="",
  745. )
  746. else:
  747. # Binary mode
  748. handle = open(handle, ioargs.mode)
  749. handles.append(handle)
  750. # Convert BytesIO or file objects passed with an encoding
  751. is_wrapped = False
  752. if not is_text and ioargs.mode == "rb" and isinstance(handle, TextIOBase):
  753. # not added to handles as it does not open/buffer resources
  754. handle = _BytesIOWrapper(
  755. handle,
  756. encoding=ioargs.encoding,
  757. )
  758. elif is_text and (
  759. compression or memory_map or _is_binary_mode(handle, ioargs.mode)
  760. ):
  761. if (
  762. not hasattr(handle, "readable")
  763. or not hasattr(handle, "writable")
  764. or not hasattr(handle, "seekable")
  765. ):
  766. handle = _IOWrapper(handle)
  767. # error: Argument 1 to "TextIOWrapper" has incompatible type
  768. # "_IOWrapper"; expected "IO[bytes]"
  769. handle = TextIOWrapper(
  770. handle, # type: ignore[arg-type]
  771. encoding=ioargs.encoding,
  772. errors=errors,
  773. newline="",
  774. )
  775. handles.append(handle)
  776. # only marked as wrapped when the caller provided a handle
  777. is_wrapped = not (
  778. isinstance(ioargs.filepath_or_buffer, str) or ioargs.should_close
  779. )
  780. if "r" in ioargs.mode and not hasattr(handle, "read"):
  781. raise TypeError(
  782. "Expected file path name or file-like object, "
  783. f"got {type(ioargs.filepath_or_buffer)} type"
  784. )
  785. handles.reverse() # close the most recently added buffer first
  786. if ioargs.should_close:
  787. assert not isinstance(ioargs.filepath_or_buffer, str)
  788. handles.append(ioargs.filepath_or_buffer)
  789. return IOHandles(
  790. # error: Argument "handle" to "IOHandles" has incompatible type
  791. # "Union[TextIOWrapper, GzipFile, BaseBuffer, typing.IO[bytes],
  792. # typing.IO[Any]]"; expected "pandas._typing.IO[Any]"
  793. handle=handle, # type: ignore[arg-type]
  794. # error: Argument "created_handles" to "IOHandles" has incompatible type
  795. # "List[BaseBuffer]"; expected "List[Union[IO[bytes], IO[str]]]"
  796. created_handles=handles, # type: ignore[arg-type]
  797. is_wrapped=is_wrapped,
  798. compression=ioargs.compression,
  799. )
  800. # error: Definition of "__enter__" in base class "IOBase" is incompatible
  801. # with definition in base class "BinaryIO"
  802. class _BufferedWriter(BytesIO, ABC): # type: ignore[misc]
  803. """
  804. Some objects do not support multiple .write() calls (TarFile and ZipFile).
  805. This wrapper writes to the underlying buffer on close.
  806. """
  807. buffer = BytesIO()
  808. @abstractmethod
  809. def write_to_buffer(self) -> None:
  810. ...
  811. def close(self) -> None:
  812. if self.closed:
  813. # already closed
  814. return
  815. if self.getbuffer().nbytes:
  816. # write to buffer
  817. self.seek(0)
  818. with self.buffer:
  819. self.write_to_buffer()
  820. else:
  821. self.buffer.close()
  822. super().close()
  823. class _BytesTarFile(_BufferedWriter):
  824. def __init__(
  825. self,
  826. name: str | None = None,
  827. mode: Literal["r", "a", "w", "x"] = "r",
  828. fileobj: ReadBuffer[bytes] | WriteBuffer[bytes] | None = None,
  829. archive_name: str | None = None,
  830. **kwargs,
  831. ) -> None:
  832. super().__init__()
  833. self.archive_name = archive_name
  834. self.name = name
  835. # error: Incompatible types in assignment (expression has type "TarFile",
  836. # base class "_BufferedWriter" defined the type as "BytesIO")
  837. self.buffer: tarfile.TarFile = tarfile.TarFile.open( # type: ignore[assignment]
  838. name=name,
  839. mode=self.extend_mode(mode),
  840. fileobj=fileobj,
  841. **kwargs,
  842. )
  843. def extend_mode(self, mode: str) -> str:
  844. mode = mode.replace("b", "")
  845. if mode != "w":
  846. return mode
  847. if self.name is not None:
  848. suffix = Path(self.name).suffix
  849. if suffix in (".gz", ".xz", ".bz2"):
  850. mode = f"{mode}:{suffix[1:]}"
  851. return mode
  852. def infer_filename(self) -> str | None:
  853. """
  854. If an explicit archive_name is not given, we still want the file inside the zip
  855. file not to be named something.tar, because that causes confusion (GH39465).
  856. """
  857. if self.name is None:
  858. return None
  859. filename = Path(self.name)
  860. if filename.suffix == ".tar":
  861. return filename.with_suffix("").name
  862. elif filename.suffix in (".tar.gz", ".tar.bz2", ".tar.xz"):
  863. return filename.with_suffix("").with_suffix("").name
  864. return filename.name
  865. def write_to_buffer(self) -> None:
  866. # TarFile needs a non-empty string
  867. archive_name = self.archive_name or self.infer_filename() or "tar"
  868. tarinfo = tarfile.TarInfo(name=archive_name)
  869. tarinfo.size = len(self.getvalue())
  870. self.buffer.addfile(tarinfo, self)
  871. class _BytesZipFile(_BufferedWriter):
  872. def __init__(
  873. self,
  874. file: FilePath | ReadBuffer[bytes] | WriteBuffer[bytes],
  875. mode: str,
  876. archive_name: str | None = None,
  877. **kwargs,
  878. ) -> None:
  879. super().__init__()
  880. mode = mode.replace("b", "")
  881. self.archive_name = archive_name
  882. kwargs.setdefault("compression", zipfile.ZIP_DEFLATED)
  883. # error: Incompatible types in assignment (expression has type "ZipFile",
  884. # base class "_BufferedWriter" defined the type as "BytesIO")
  885. self.buffer: zipfile.ZipFile = zipfile.ZipFile( # type: ignore[assignment]
  886. file, mode, **kwargs
  887. )
  888. def infer_filename(self) -> str | None:
  889. """
  890. If an explicit archive_name is not given, we still want the file inside the zip
  891. file not to be named something.zip, because that causes confusion (GH39465).
  892. """
  893. if isinstance(self.buffer.filename, (os.PathLike, str)):
  894. filename = Path(self.buffer.filename)
  895. if filename.suffix == ".zip":
  896. return filename.with_suffix("").name
  897. return filename.name
  898. return None
  899. def write_to_buffer(self) -> None:
  900. # ZipFile needs a non-empty string
  901. archive_name = self.archive_name or self.infer_filename() or "zip"
  902. self.buffer.writestr(archive_name, self.getvalue())
  903. class _IOWrapper:
  904. # TextIOWrapper is overly strict: it request that the buffer has seekable, readable,
  905. # and writable. If we have a read-only buffer, we shouldn't need writable and vice
  906. # versa. Some buffers, are seek/read/writ-able but they do not have the "-able"
  907. # methods, e.g., tempfile.SpooledTemporaryFile.
  908. # If a buffer does not have the above "-able" methods, we simple assume they are
  909. # seek/read/writ-able.
  910. def __init__(self, buffer: BaseBuffer) -> None:
  911. self.buffer = buffer
  912. def __getattr__(self, name: str):
  913. return getattr(self.buffer, name)
  914. def readable(self) -> bool:
  915. if hasattr(self.buffer, "readable"):
  916. return self.buffer.readable()
  917. return True
  918. def seekable(self) -> bool:
  919. if hasattr(self.buffer, "seekable"):
  920. return self.buffer.seekable()
  921. return True
  922. def writable(self) -> bool:
  923. if hasattr(self.buffer, "writable"):
  924. return self.buffer.writable()
  925. return True
  926. class _BytesIOWrapper:
  927. # Wrapper that wraps a StringIO buffer and reads bytes from it
  928. # Created for compat with pyarrow read_csv
  929. def __init__(self, buffer: StringIO | TextIOBase, encoding: str = "utf-8") -> None:
  930. self.buffer = buffer
  931. self.encoding = encoding
  932. # Because a character can be represented by more than 1 byte,
  933. # it is possible that reading will produce more bytes than n
  934. # We store the extra bytes in this overflow variable, and append the
  935. # overflow to the front of the bytestring the next time reading is performed
  936. self.overflow = b""
  937. def __getattr__(self, attr: str):
  938. return getattr(self.buffer, attr)
  939. def read(self, n: int | None = -1) -> bytes:
  940. assert self.buffer is not None
  941. bytestring = self.buffer.read(n).encode(self.encoding)
  942. # When n=-1/n greater than remaining bytes: Read entire file/rest of file
  943. combined_bytestring = self.overflow + bytestring
  944. if n is None or n < 0 or n >= len(combined_bytestring):
  945. self.overflow = b""
  946. return combined_bytestring
  947. else:
  948. to_return = combined_bytestring[:n]
  949. self.overflow = combined_bytestring[n:]
  950. return to_return
  951. def _maybe_memory_map(
  952. handle: str | BaseBuffer, memory_map: bool
  953. ) -> tuple[str | BaseBuffer, bool, list[BaseBuffer]]:
  954. """Try to memory map file/buffer."""
  955. handles: list[BaseBuffer] = []
  956. memory_map &= hasattr(handle, "fileno") or isinstance(handle, str)
  957. if not memory_map:
  958. return handle, memory_map, handles
  959. # mmap used by only read_csv
  960. handle = cast(ReadCsvBuffer, handle)
  961. # need to open the file first
  962. if isinstance(handle, str):
  963. handle = open(handle, "rb")
  964. handles.append(handle)
  965. try:
  966. # open mmap and adds *-able
  967. # error: Argument 1 to "_IOWrapper" has incompatible type "mmap";
  968. # expected "BaseBuffer"
  969. wrapped = _IOWrapper(
  970. mmap.mmap(
  971. handle.fileno(), 0, access=mmap.ACCESS_READ # type: ignore[arg-type]
  972. )
  973. )
  974. finally:
  975. for handle in reversed(handles):
  976. # error: "BaseBuffer" has no attribute "close"
  977. handle.close() # type: ignore[attr-defined]
  978. return wrapped, memory_map, [wrapped]
  979. def file_exists(filepath_or_buffer: FilePath | BaseBuffer) -> bool:
  980. """Test whether file exists."""
  981. exists = False
  982. filepath_or_buffer = stringify_path(filepath_or_buffer)
  983. if not isinstance(filepath_or_buffer, str):
  984. return exists
  985. try:
  986. exists = os.path.exists(filepath_or_buffer)
  987. # gh-5874: if the filepath is too long will raise here
  988. except (TypeError, ValueError):
  989. pass
  990. return exists
  991. def _is_binary_mode(handle: FilePath | BaseBuffer, mode: str) -> bool:
  992. """Whether the handle is opened in binary mode"""
  993. # specified by user
  994. if "t" in mode or "b" in mode:
  995. return "b" in mode
  996. # exceptions
  997. text_classes = (
  998. # classes that expect string but have 'b' in mode
  999. codecs.StreamWriter,
  1000. codecs.StreamReader,
  1001. codecs.StreamReaderWriter,
  1002. )
  1003. if issubclass(type(handle), text_classes):
  1004. return False
  1005. return isinstance(handle, _get_binary_io_classes()) or "b" in getattr(
  1006. handle, "mode", mode
  1007. )
  1008. @functools.lru_cache
  1009. def _get_binary_io_classes() -> tuple[type, ...]:
  1010. """IO classes that that expect bytes"""
  1011. binary_classes: tuple[type, ...] = (BufferedIOBase, RawIOBase)
  1012. # python-zstandard doesn't use any of the builtin base classes; instead we
  1013. # have to use the `zstd.ZstdDecompressionReader` class for isinstance checks.
  1014. # Unfortunately `zstd.ZstdDecompressionReader` isn't exposed by python-zstandard
  1015. # so we have to get it from a `zstd.ZstdDecompressor` instance.
  1016. # See also https://github.com/indygreg/python-zstandard/pull/165.
  1017. zstd = import_optional_dependency("zstandard", errors="ignore")
  1018. if zstd is not None:
  1019. with zstd.ZstdDecompressor().stream_reader(b"") as reader:
  1020. binary_classes += (type(reader),)
  1021. return binary_classes
  1022. def is_potential_multi_index(
  1023. columns: Sequence[Hashable] | MultiIndex,
  1024. index_col: bool | Sequence[int] | None = None,
  1025. ) -> bool:
  1026. """
  1027. Check whether or not the `columns` parameter
  1028. could be converted into a MultiIndex.
  1029. Parameters
  1030. ----------
  1031. columns : array-like
  1032. Object which may or may not be convertible into a MultiIndex
  1033. index_col : None, bool or list, optional
  1034. Column or columns to use as the (possibly hierarchical) index
  1035. Returns
  1036. -------
  1037. bool : Whether or not columns could become a MultiIndex
  1038. """
  1039. if index_col is None or isinstance(index_col, bool):
  1040. index_col = []
  1041. return bool(
  1042. len(columns)
  1043. and not isinstance(columns, ABCMultiIndex)
  1044. and all(isinstance(c, tuple) for c in columns if c not in list(index_col))
  1045. )
  1046. def dedup_names(
  1047. names: Sequence[Hashable], is_potential_multiindex: bool
  1048. ) -> Sequence[Hashable]:
  1049. """
  1050. Rename column names if duplicates exist.
  1051. Currently the renaming is done by appending a period and an autonumeric,
  1052. but a custom pattern may be supported in the future.
  1053. Examples
  1054. --------
  1055. >>> dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False)
  1056. ['x', 'y', 'x.1', 'x.2']
  1057. """
  1058. names = list(names) # so we can index
  1059. counts: DefaultDict[Hashable, int] = defaultdict(int)
  1060. for i, col in enumerate(names):
  1061. cur_count = counts[col]
  1062. while cur_count > 0:
  1063. counts[col] = cur_count + 1
  1064. if is_potential_multiindex:
  1065. # for mypy
  1066. assert isinstance(col, tuple)
  1067. col = col[:-1] + (f"{col[-1]}.{cur_count}",)
  1068. else:
  1069. col = f"{col}.{cur_count}"
  1070. cur_count = counts[col]
  1071. names[i] = col
  1072. counts[col] = cur_count + 1
  1073. return names