Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/common.py: 22%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Common I/O API utilities"""
3from __future__ import annotations
5from abc import (
6 ABC,
7 abstractmethod,
8)
9import codecs
10from collections import defaultdict
11from collections.abc import (
12 Hashable,
13 Mapping,
14 Sequence,
15)
16import dataclasses
17import functools
18import gzip
19from io import (
20 BufferedIOBase,
21 BytesIO,
22 IOBase,
23 RawIOBase,
24 StringIO,
25 TextIOBase,
26 TextIOWrapper,
27 UnsupportedOperation,
28)
29import mmap
30import os
31from pathlib import Path
32import re
33import tarfile
34from typing import (
35 IO,
36 TYPE_CHECKING,
37 Any,
38 AnyStr,
39 Generic,
40 Literal,
41 TypeVar,
42 cast,
43 overload,
44)
45from urllib.parse import (
46 urljoin,
47 urlparse as parse_url,
48 uses_netloc,
49 uses_params,
50 uses_relative,
51)
52import warnings
53import zipfile
55from pandas._typing import (
56 BaseBuffer,
57 ReadCsvBuffer,
58)
59from pandas.compat._optional import import_optional_dependency
60from pandas.util._exceptions import find_stack_level
62from pandas.core.dtypes.common import (
63 is_bool,
64 is_file_like,
65 is_integer,
66 is_list_like,
67)
68from pandas.core.dtypes.generic import ABCMultiIndex
70_VALID_URLS = set(uses_relative + uses_netloc + uses_params)
71_VALID_URLS.discard("")
72_FSSPEC_URL_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9+\-+.]*(::[A-Za-z0-9+\-+.]+)*://")
74BaseBufferT = TypeVar("BaseBufferT", bound=BaseBuffer)
77if TYPE_CHECKING:
78 from types import TracebackType
80 from pandas._typing import (
81 CompressionDict,
82 CompressionOptions,
83 DtypeArg,
84 FilePath,
85 ReadBuffer,
86 StorageOptions,
87 WriteBuffer,
88 )
90 from pandas import MultiIndex
93@dataclasses.dataclass
94class IOArgs:
95 """
96 Return value of io/common.py:_get_filepath_or_buffer.
97 """
99 filepath_or_buffer: str | BaseBuffer
100 encoding: str
101 mode: str
102 compression: CompressionDict
103 should_close: bool = False
104 close_handles: list[Any] = dataclasses.field(default_factory=list)
107@dataclasses.dataclass
108class IOHandles(Generic[AnyStr]):
109 """
110 Return value of io/common.py:get_handle
112 Can be used as a context manager.
114 This is used to easily close created buffers and to handle corner cases when
115 TextIOWrapper is inserted.
117 handle: The file handle to be used.
118 created_handles: All file handles that are created by get_handle
119 is_wrapped: Whether a TextIOWrapper needs to be detached.
120 """
122 # handle might not implement the IO-interface
123 handle: IO[AnyStr]
124 compression: CompressionDict
125 created_handles: list[IO[bytes] | IO[str]] = dataclasses.field(default_factory=list)
126 is_wrapped: bool = False
128 def close(self) -> None:
129 """
130 Close all created buffers.
132 Note: If a TextIOWrapper was inserted, it is flushed and detached to
133 avoid closing the potentially user-created buffer.
134 """
135 if self.is_wrapped:
136 assert isinstance(self.handle, TextIOWrapper)
137 self.handle.flush()
138 self.handle.detach()
139 self.created_handles.remove(self.handle)
140 for handle in self.created_handles:
141 handle.close()
142 self.created_handles = []
143 self.is_wrapped = False
145 def __enter__(self) -> IOHandles[AnyStr]:
146 return self
148 def __exit__(
149 self,
150 exc_type: type[BaseException] | None,
151 exc_value: BaseException | None,
152 traceback: TracebackType | None,
153 ) -> None:
154 self.close()
157def is_url(url: object) -> bool:
158 """
159 Check to see if a URL has a valid protocol.
161 Parameters
162 ----------
163 url : str or unicode
165 Returns
166 -------
167 isurl : bool
168 If `url` has a valid protocol return True otherwise False.
169 """
170 if not isinstance(url, str):
171 return False
172 return parse_url(url).scheme in _VALID_URLS
175@overload
176def _expand_user(filepath_or_buffer: str) -> str: ...
179@overload
180def _expand_user(filepath_or_buffer: BaseBufferT) -> BaseBufferT: ...
183def _expand_user(filepath_or_buffer: str | BaseBufferT) -> str | BaseBufferT:
184 """
185 Return the argument with an initial component of ~ or ~user
186 replaced by that user's home directory.
188 Parameters
189 ----------
190 filepath_or_buffer : object to be converted if possible
192 Returns
193 -------
194 expanded_filepath_or_buffer : an expanded filepath or the
195 input if not expandable
196 """
197 if isinstance(filepath_or_buffer, str):
198 return os.path.expanduser(filepath_or_buffer)
199 return filepath_or_buffer
202def validate_header_arg(header: object) -> None:
203 if header is None:
204 return
205 if is_integer(header):
206 header = cast("int", header)
207 if header < 0:
208 # GH 27779
209 raise ValueError(
210 "Passing negative integer to header is invalid. "
211 "For no header, use header=None instead"
212 )
213 return
214 if is_list_like(header, allow_sets=False):
215 header = cast("Sequence[int]", header)
216 if not all(map(is_integer, header)):
217 raise ValueError("header must be integer or list of integers")
218 if any(i < 0 for i in header):
219 raise ValueError("cannot specify multi-index header with negative integers")
220 return
221 if is_bool(header):
222 raise TypeError(
223 "Passing a bool to header is invalid. Use header=None for no header or "
224 "header=int or list-like of ints to specify "
225 "the row(s) making up the column names"
226 )
227 # GH 16338
228 raise ValueError("header must be integer or list of integers")
231@overload
232def stringify_path(
233 filepath_or_buffer: FilePath, convert_file_like: bool = ...
234) -> str: ...
237@overload
238def stringify_path(
239 filepath_or_buffer: BaseBufferT, convert_file_like: bool = ...
240) -> BaseBufferT: ...
243def stringify_path(
244 filepath_or_buffer: FilePath | BaseBufferT,
245 convert_file_like: bool = False,
246) -> str | BaseBufferT:
247 """
248 Attempt to convert a path-like object to a string.
250 Parameters
251 ----------
252 filepath_or_buffer : object to be converted
254 Returns
255 -------
256 str_filepath_or_buffer : maybe a string version of the object
258 Notes
259 -----
260 Objects supporting the fspath protocol are coerced
261 according to its __fspath__ method.
263 Any other object is passed through unchanged, which includes bytes,
264 strings, buffers, or anything else that's not even path-like.
265 """
266 if (
267 not convert_file_like
268 and is_file_like(filepath_or_buffer)
269 and hasattr(filepath_or_buffer, "__iter__")
270 ):
271 # GH 38125: some fsspec objects implement os.PathLike but have already opened a
272 # file. This prevents opening the file a second time. infer_compression calls
273 # this function with convert_file_like=True to infer the compression.
274 # Requiring __iter__ keeps py.path.local, whose read/write are convenience
275 # methods, on the path branch (GH#71769).
276 return cast("BaseBufferT", filepath_or_buffer)
278 if isinstance(filepath_or_buffer, os.PathLike):
279 filepath_or_buffer = filepath_or_buffer.__fspath__()
280 return _expand_user(filepath_or_buffer)
283def urlopen(*args: Any, **kwargs: Any) -> Any:
284 """
285 Lazy-import wrapper for stdlib urlopen, as that imports a big chunk of
286 the stdlib.
287 """
288 import urllib.request
290 return urllib.request.urlopen(*args, **kwargs) # noqa: TID251
293def is_fsspec_url(url: FilePath | BaseBuffer) -> bool:
294 """
295 Returns true if the given URL looks like
296 something fsspec can handle
297 """
298 return (
299 isinstance(url, str)
300 and bool(_FSSPEC_URL_PATTERN.match(url))
301 and not url.startswith(("http://", "https://"))
302 )
305def _get_filepath_or_buffer(
306 filepath_or_buffer: FilePath | BaseBuffer,
307 encoding: str = "utf-8",
308 compression: CompressionOptions | None = None,
309 mode: str = "r",
310 storage_options: StorageOptions | None = None,
311) -> IOArgs:
312 """
313 If the filepath_or_buffer is a url, translate and return the buffer.
314 Otherwise passthrough.
316 Parameters
317 ----------
318 filepath_or_buffer : a url, filepath (str or pathlib.Path),
319 or buffer
321 compression : str or dict, default 'infer'
322 For on-the-fly compression of the output data. If 'infer' and
323 'filepath_or_buffer' is path-like, then detect compression from the
324 following extensions: '.gz',
325 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
326 (otherwise no compression).
327 Set to ``None`` for no compression.
328 Can also be a dict with key ``'method'`` set
329 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
330 and other key-value pairs are forwarded to
331 ``zipfile.ZipFile``, ``gzip.GzipFile``,
332 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
333 ``tarfile.TarFile``, respectively.
334 As an example, the following could be passed for faster compression and to
335 create a reproducible gzip archive:
336 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
338 encoding : the encoding to use to decode bytes, default is 'utf-8'
339 mode : str, optional
341 storage_options : dict, optional
342 Extra options that make sense for a particular storage connection, e.g.
343 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs
344 are forwarded to ``urllib.request.Request`` as header options. For other
345 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are
346 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more
347 details, and for more examples on storage options refer `here
348 <https://pandas.pydata.org/docs/user_guide/io.html?
349 highlight=storage_options#reading-writing-remote-files>`_.
352 Returns the dataclass IOArgs.
353 """
354 filepath_or_buffer = stringify_path(filepath_or_buffer)
356 # handle compression dict
357 compression_method, compression = get_compression_method(compression)
358 compression_method = infer_compression(filepath_or_buffer, compression_method)
360 # GH21227 internal compression is not used for non-binary handles.
361 if compression_method and hasattr(filepath_or_buffer, "write") and "b" not in mode:
362 warnings.warn(
363 "compression has no effect when passing a non-binary object as input.",
364 RuntimeWarning,
365 stacklevel=find_stack_level(),
366 )
367 compression_method = None
369 compression = dict(compression, method=compression_method)
371 # bz2 and xz do not write the byte order mark for utf-16 and utf-32
372 # print a warning when writing such files
373 if (
374 "w" in mode
375 and compression_method in ["bz2", "xz"]
376 and encoding in ["utf-16", "utf-32"]
377 ):
378 warnings.warn(
379 f"{compression_method} will not write the byte order mark for {encoding}",
380 UnicodeWarning,
381 stacklevel=find_stack_level(),
382 )
384 if "a" in mode and compression_method in ["zip", "tar"]:
385 # GH56778
386 warnings.warn(
387 "zip and tar do not support mode 'a' properly. "
388 "This combination will result in multiple files with same name "
389 "being added to the archive.",
390 RuntimeWarning,
391 stacklevel=find_stack_level(),
392 )
394 # Use binary mode when converting path-like objects to file-like objects (fsspec)
395 # except when text mode is explicitly requested. The original mode is returned if
396 # fsspec is not used.
397 fsspec_mode = mode
398 if "t" not in fsspec_mode and "b" not in fsspec_mode:
399 fsspec_mode += "b"
401 if isinstance(filepath_or_buffer, str) and is_url(filepath_or_buffer):
402 # TODO: fsspec can also handle HTTP via requests, but leaving this
403 # unchanged. using fsspec appears to break the ability to infer if the
404 # server responded with gzipped data
405 storage_options = storage_options or {}
407 # waiting until now for importing to match intended lazy logic of
408 # urlopen function defined elsewhere in this module
409 import urllib.request
411 # assuming storage_options is to be interpreted as headers
412 req_info = urllib.request.Request(filepath_or_buffer, headers=storage_options)
413 with urlopen(req_info) as req:
414 content_encoding = req.headers.get("Content-Encoding", None)
415 if content_encoding == "gzip":
416 # Override compression based on Content-Encoding header
417 compression = {"method": "gzip"}
418 reader = BytesIO(req.read())
419 return IOArgs(
420 filepath_or_buffer=reader,
421 encoding=encoding,
422 compression=compression,
423 should_close=True,
424 mode=fsspec_mode,
425 )
427 if is_fsspec_url(filepath_or_buffer):
428 assert isinstance(
429 filepath_or_buffer, str
430 ) # just to appease mypy for this branch
431 # two special-case s3-like protocols; these have special meaning in Hadoop,
432 # but are equivalent to just "s3" from fsspec's point of view
433 # cc #11071
434 if filepath_or_buffer.startswith("s3a://"):
435 filepath_or_buffer = filepath_or_buffer.replace("s3a://", "s3://")
436 if filepath_or_buffer.startswith("s3n://"):
437 filepath_or_buffer = filepath_or_buffer.replace("s3n://", "s3://")
438 fsspec = import_optional_dependency("fsspec")
440 # If botocore is installed we fallback to reading with anon=True
441 # to allow reads from public buckets
442 err_types_to_retry_with_anon: list[Any] = []
443 try:
444 import_optional_dependency("botocore")
445 from botocore.exceptions import (
446 ClientError,
447 NoCredentialsError,
448 )
450 err_types_to_retry_with_anon = [
451 ClientError,
452 NoCredentialsError,
453 PermissionError,
454 ]
455 except ImportError:
456 pass
458 try:
459 open_file = fsspec.open(
460 filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
461 )
462 file_obj = open_file.open()
463 # GH 34626 Reads from Public Buckets without Credentials needs anon=True
464 except tuple(err_types_to_retry_with_anon):
465 if storage_options is None:
466 storage_options = {"anon": True}
467 else:
468 # don't mutate user input.
469 storage_options = dict(storage_options)
470 storage_options["anon"] = True
471 open_file = fsspec.open(
472 filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
473 )
474 file_obj = open_file.open()
476 return IOArgs(
477 filepath_or_buffer=file_obj,
478 encoding=encoding,
479 compression=compression,
480 close_handles=[open_file],
481 should_close=True,
482 mode=fsspec_mode,
483 )
484 elif storage_options:
485 raise ValueError(
486 "storage_options passed with file object or non-fsspec file path"
487 )
489 if isinstance(filepath_or_buffer, (str, bytes, mmap.mmap)):
490 return IOArgs(
491 filepath_or_buffer=_expand_user(filepath_or_buffer),
492 encoding=encoding,
493 compression=compression,
494 should_close=False,
495 mode=mode,
496 )
498 if not (
499 hasattr(filepath_or_buffer, "read") or hasattr(filepath_or_buffer, "write")
500 ):
501 msg = f"Invalid file path or buffer object type: {type(filepath_or_buffer)}"
502 raise ValueError(msg)
504 return IOArgs(
505 filepath_or_buffer=filepath_or_buffer,
506 encoding=encoding,
507 compression=compression,
508 should_close=False,
509 mode=mode,
510 )
513def file_path_to_url(path: str) -> str:
514 """
515 converts an absolute native path to a FILE URL.
517 Parameters
518 ----------
519 path : a path in native format
521 Returns
522 -------
523 a valid FILE URL
524 """
525 # lazify expensive import (~30ms)
526 from urllib.request import pathname2url
528 return urljoin("file:", pathname2url(path))
531extension_to_compression = {
532 ".tar": "tar",
533 ".tar.gz": "tar",
534 ".tar.bz2": "tar",
535 ".tar.xz": "tar",
536 ".gz": "gzip",
537 ".bz2": "bz2",
538 ".zip": "zip",
539 ".xz": "xz",
540 ".zst": "zstd",
541}
542_supported_compressions = set(extension_to_compression.values())
545def get_compression_method(
546 compression: CompressionOptions,
547) -> tuple[str | None, CompressionDict]:
548 """
549 Simplifies a compression argument to a compression method string and
550 a mapping containing additional arguments.
552 Parameters
553 ----------
554 compression : str or mapping
555 If string, specifies the compression method. If mapping, value at key
556 'method' specifies compression method.
558 Returns
559 -------
560 tuple of ({compression method}, Optional[str]
561 {compression arguments}, Dict[str, Any])
563 Raises
564 ------
565 ValueError on mapping missing 'method' key
566 """
567 compression_method: str | None
568 if isinstance(compression, Mapping):
569 compression_args = dict(compression)
570 try:
571 compression_method = compression_args.pop("method")
572 except KeyError as err:
573 raise ValueError("If mapping, compression must have key 'method'") from err
574 else:
575 compression_args = {}
576 compression_method = compression
577 return compression_method, compression_args
580def infer_compression(
581 filepath_or_buffer: FilePath | BaseBuffer, compression: str | None
582) -> str | None:
583 """
584 Get the compression method for filepath_or_buffer. If compression='infer',
585 the inferred compression method is returned. Otherwise, the input
586 compression method is returned unchanged, unless it's invalid, in which
587 case an error is raised.
589 Parameters
590 ----------
591 filepath_or_buffer : str or file handle
592 File path or object.
594 compression : str or dict, default 'infer'
595 For on-the-fly compression of the output data. If 'infer' and
596 'filepath_or_buffer' is path-like, then detect compression from the
597 following extensions: '.gz',
598 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
599 (otherwise no compression).
600 Set to ``None`` for no compression.
601 Can also be a dict with key ``'method'`` set
602 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
603 and other key-value pairs are forwarded to
604 ``zipfile.ZipFile``, ``gzip.GzipFile``,
605 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
606 ``tarfile.TarFile``, respectively.
607 As an example, the following could be passed for faster compression and to
608 create a reproducible gzip archive:
609 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
611 Returns
612 -------
613 string or None
615 Raises
616 ------
617 ValueError on invalid compression specified.
618 """
619 if compression is None:
620 return None
622 # Infer compression
623 if compression == "infer":
624 # Convert all path types (e.g. pathlib.Path) to strings
625 if isinstance(filepath_or_buffer, str) and "::" in filepath_or_buffer:
626 # chained URLs contain ::
627 filepath_or_buffer = filepath_or_buffer.split("::")[0]
628 filepath_or_buffer = stringify_path(filepath_or_buffer, convert_file_like=True)
629 if not isinstance(filepath_or_buffer, str):
630 # Cannot infer compression of a buffer, assume no compression
631 return None
633 # Infer compression from the filename/URL extension
634 for extension, compression in extension_to_compression.items():
635 if filepath_or_buffer.lower().endswith(extension):
636 return compression
637 return None
639 # Compression has been specified. Check that it's valid
640 if compression in _supported_compressions:
641 return compression
643 valid = ["infer", None, *sorted(_supported_compressions)]
644 msg = (
645 f"Unrecognized compression type: {compression}\n"
646 f"Valid compression types are {valid}"
647 )
648 raise ValueError(msg)
651def check_parent_directory(path: Path | str) -> None:
652 """
653 Check if parent directory of a file exists, raise OSError if it does not
655 Parameters
656 ----------
657 path: Path or str
658 Path to check parent directory of
659 """
660 parent = Path(path).parent
661 if not parent.is_dir():
662 raise OSError(rf"Cannot save file into a non-existent directory: '{parent}'")
665@overload
666def get_handle(
667 path_or_buf: FilePath | BaseBuffer,
668 mode: str,
669 *,
670 encoding: str | None = ...,
671 compression: CompressionOptions = ...,
672 memory_map: bool = ...,
673 is_text: Literal[False],
674 errors: str | None = ...,
675 storage_options: StorageOptions = ...,
676) -> IOHandles[bytes]: ...
679@overload
680def get_handle(
681 path_or_buf: FilePath | BaseBuffer,
682 mode: str,
683 *,
684 encoding: str | None = ...,
685 compression: CompressionOptions = ...,
686 memory_map: bool = ...,
687 is_text: Literal[True] = ...,
688 errors: str | None = ...,
689 storage_options: StorageOptions = ...,
690) -> IOHandles[str]: ...
693@overload
694def get_handle(
695 path_or_buf: FilePath | BaseBuffer,
696 mode: str,
697 *,
698 encoding: str | None = ...,
699 compression: CompressionOptions = ...,
700 memory_map: bool = ...,
701 is_text: bool = ...,
702 errors: str | None = ...,
703 storage_options: StorageOptions = ...,
704) -> IOHandles[str] | IOHandles[bytes]: ...
707def get_handle(
708 path_or_buf: FilePath | BaseBuffer,
709 mode: str,
710 *,
711 encoding: str | None = None,
712 compression: CompressionOptions | None = None,
713 memory_map: bool = False,
714 is_text: bool = True,
715 errors: str | None = None,
716 storage_options: StorageOptions | None = None,
717) -> IOHandles[str] | IOHandles[bytes]:
718 """
719 Get file handle for given path/buffer and mode.
721 Parameters
722 ----------
723 path_or_buf : str or file handle
724 File path or object.
725 mode : str
726 Mode to open path_or_buf with.
727 encoding : str or None
728 Encoding to use.
729 compression : str or dict, default 'infer'
730 For on-the-fly compression of the output data. If 'infer' and 'path_or_buf'
731 is path-like, then detect compression from the following extensions: '.gz',
732 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
733 (otherwise no compression).
734 Set to ``None`` for no compression.
735 Can also be a dict with key ``'method'`` set
736 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
737 and other key-value pairs are forwarded to
738 ``zipfile.ZipFile``, ``gzip.GzipFile``,
739 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
740 ``tarfile.TarFile``, respectively.
741 As an example, the following could be passed for faster compression and to
742 create a reproducible gzip archive:
743 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
745 May be a dict with key 'method' as compression mode
746 and other keys as compression options if compression
747 mode is 'zip'.
749 Passing compression options as keys in dict is
750 supported for compression modes 'gzip', 'bz2', 'zstd' and 'zip'.
752 memory_map : bool, default False
753 See parsers._parser_params for more information. Only used by read_csv.
754 is_text : bool, default True
755 Whether the type of the content passed to the file/buffer is string or
756 bytes. This is not the same as `"b" not in mode`. If a string content is
757 passed to a binary file/buffer, a wrapper is inserted.
758 errors : str, default 'strict'
759 Specifies how encoding and decoding errors are to be handled.
760 See the errors argument for :func:`open` for a full list
761 of options.
762 storage_options: StorageOptions = None
763 Passed to _get_filepath_or_buffer
765 Returns the dataclass IOHandles
766 """
767 # Windows does not default to utf-8. Set to utf-8 for a consistent behavior
768 encoding = encoding or "utf-8"
770 errors = errors or "strict"
772 # read_csv does not know whether the buffer is opened in binary/text mode
773 if _is_binary_mode(path_or_buf, mode) and "b" not in mode:
774 mode += "b"
776 # validate encoding and errors
777 codecs.lookup(encoding)
778 if isinstance(errors, str):
779 codecs.lookup_error(errors)
781 # open URLs
782 ioargs = _get_filepath_or_buffer(
783 path_or_buf,
784 encoding=encoding,
785 compression=compression,
786 mode=mode,
787 storage_options=storage_options,
788 )
790 handle = ioargs.filepath_or_buffer
791 handles: list[BaseBuffer] = []
792 completed = False
794 try:
795 # memory mapping needs to be the first step
796 # only used for read_csv
797 handle, memory_map, handles = _maybe_memory_map(handle, memory_map)
799 is_path = isinstance(handle, str)
800 compression_args = dict(ioargs.compression)
801 compression = compression_args.pop("method")
803 # Only for write methods
804 if "r" not in mode and is_path:
805 check_parent_directory(str(handle))
807 if compression:
808 if compression != "zstd":
809 # compression libraries do not like an explicit text-mode
810 ioargs.mode = ioargs.mode.replace("t", "")
811 elif compression == "zstd" and "b" not in ioargs.mode:
812 # python-zstandard defaults to text mode, but we always expect
813 # compression libraries to use binary mode.
814 ioargs.mode += "b"
816 # GZ Compression
817 if compression == "gzip":
818 if isinstance(handle, str):
819 handle = gzip.GzipFile(
820 filename=handle,
821 mode=ioargs.mode,
822 **compression_args,
823 )
824 else:
825 handle = gzip.GzipFile(
826 # No overload variant of "GzipFile" matches argument types
827 # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
828 fileobj=handle, # type: ignore[call-overload]
829 mode=ioargs.mode,
830 **compression_args,
831 )
833 # BZ Compression
834 elif compression == "bz2":
835 import bz2
837 # Overload of "BZ2File" to handle pickle protocol 5
838 # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
839 handle = bz2.BZ2File( # type: ignore[call-overload]
840 handle,
841 mode=ioargs.mode,
842 **compression_args,
843 )
845 # ZIP Compression
846 elif compression == "zip":
847 # error: Argument 1 to "_BytesZipFile" has incompatible type
848 # "Union[str, BaseBuffer]"; expected "Union[Union[str, PathLike[str]],
849 # ReadBuffer[bytes], WriteBuffer[bytes]]"
850 handle = _BytesZipFile(
851 handle, # type: ignore[arg-type]
852 ioargs.mode,
853 **compression_args,
854 )
855 if handle.buffer.mode == "r":
856 handles.append(handle)
857 zip_names = handle.buffer.namelist()
858 if len(zip_names) == 1:
859 # Read the entire zip entry into memory rather than
860 # returning a streaming ZipExtFile via .open(). On
861 # Python <3.12, ZipExtFile has poor buffering that
862 # causes O(n²) performance with pickle.load() and
863 # other consumers that issue many small reads
864 # (GH#59279). Can be reverted to .open() once the
865 # minimum Python version is 3.12.
866 handle = BytesIO(handle.buffer.read(zip_names.pop()))
867 elif not zip_names:
868 raise ValueError(f"Zero files found in ZIP file {path_or_buf}")
869 else:
870 raise ValueError(
871 "Multiple files found in ZIP file. "
872 f"Only one file per ZIP: {zip_names}"
873 )
875 # TAR Encoding
876 elif compression == "tar":
877 compression_args.setdefault("mode", ioargs.mode)
878 if isinstance(handle, str):
879 handle = _BytesTarFile(name=handle, **compression_args)
880 else:
881 # error: Argument "fileobj" to "_BytesTarFile" has incompatible
882 # type "BaseBuffer"; expected "Union[ReadBuffer[bytes],
883 # WriteBuffer[bytes], None]"
884 handle = _BytesTarFile(
885 fileobj=handle, # type: ignore[arg-type]
886 **compression_args,
887 )
888 assert isinstance(handle, _BytesTarFile)
889 if "r" in handle.buffer.mode:
890 handles.append(handle)
891 files = handle.buffer.getnames()
892 if len(files) == 1:
893 file = handle.buffer.extractfile(files[0])
894 assert file is not None
895 handle = file
896 elif not files:
897 raise ValueError(
898 f"Zero files found in TAR archive {path_or_buf}"
899 )
900 else:
901 raise ValueError(
902 "Multiple files found in TAR archive. "
903 f"Only one file per TAR archive: {files}"
904 )
906 # XZ Compression
907 elif compression == "xz":
908 # error: Argument 1 to "LZMAFile" has incompatible type
909 # "Union[str, BaseBuffer]"; expected "Optional[Union[Union[str,
910 # bytes, PathLike[str], PathLike[bytes]], IO[bytes]], None]"
911 import lzma
913 handle = lzma.LZMAFile(
914 handle, # type: ignore[arg-type]
915 ioargs.mode,
916 **compression_args,
917 )
919 # Zstd Compression
920 elif compression == "zstd":
921 zstd = import_optional_dependency("zstandard")
922 if "r" in ioargs.mode:
923 open_args = {"dctx": zstd.ZstdDecompressor(**compression_args)}
924 else:
925 open_args = {"cctx": zstd.ZstdCompressor(**compression_args)}
926 handle = zstd.open(
927 handle,
928 mode=ioargs.mode,
929 **open_args,
930 )
932 # Unrecognized Compression
933 else:
934 msg = f"Unrecognized compression type: {compression}"
935 raise ValueError(msg)
937 assert not isinstance(handle, str)
938 handles.append(handle)
940 elif isinstance(handle, str):
941 # Check whether the filename is to be opened in binary mode.
942 # Binary mode does not support 'encoding' and 'newline'.
943 if ioargs.encoding and "b" not in ioargs.mode:
944 # Encoding
945 handle = open(
946 handle,
947 ioargs.mode,
948 encoding=ioargs.encoding,
949 errors=errors,
950 newline="",
951 )
952 else:
953 # Binary mode
954 handle = open(handle, ioargs.mode)
955 handles.append(handle)
957 # Convert BytesIO or file objects passed with an encoding
958 is_wrapped = False
959 if not is_text and ioargs.mode == "rb" and isinstance(handle, TextIOBase):
960 # not added to handles as it does not open/buffer resources
961 handle = _BytesIOWrapper(
962 handle,
963 encoding=ioargs.encoding,
964 errors=errors,
965 )
966 elif is_text and (
967 compression or memory_map or _is_binary_mode(handle, ioargs.mode)
968 ):
969 if (
970 not hasattr(handle, "readable")
971 or not hasattr(handle, "writable")
972 or not hasattr(handle, "seekable")
973 ):
974 handle = _IOWrapper(handle)
975 # error: Value of type variable "_BufferT_co" of "TextIOWrapper" cannot
976 # be "_IOWrapper | BaseBuffer" [type-var]
977 handle = TextIOWrapper(
978 handle, # type: ignore[type-var]
979 encoding=ioargs.encoding,
980 errors=errors,
981 newline="",
982 )
983 handles.append(handle)
984 # only marked as wrapped when the caller provided a handle
985 is_wrapped = not (
986 isinstance(ioargs.filepath_or_buffer, str) or ioargs.should_close
987 )
989 if "r" in ioargs.mode and not hasattr(handle, "read"):
990 raise TypeError(
991 "Expected file path name or file-like object, "
992 f"got {type(ioargs.filepath_or_buffer)} type"
993 )
995 completed = True
996 finally:
997 if not completed:
998 # Do not leak the handles that were opened before the failure; the
999 # caller never receives an IOHandles to close them with (GH#58131)
1000 for created_handle in reversed(handles):
1001 created_handle.close() # type: ignore[attr-defined]
1002 if ioargs.should_close:
1003 assert not isinstance(ioargs.filepath_or_buffer, str)
1004 ioargs.filepath_or_buffer.close() # type: ignore[attr-defined]
1005 for close_handle in ioargs.close_handles:
1006 close_handle.close()
1008 handles.reverse() # close the most recently added buffer first
1009 if ioargs.should_close:
1010 assert not isinstance(ioargs.filepath_or_buffer, str)
1011 handles.append(ioargs.filepath_or_buffer)
1012 handles.extend(ioargs.close_handles)
1014 return IOHandles(
1015 # error: Argument "handle" to "IOHandles" has incompatible type
1016 # "Union[TextIOWrapper, GzipFile, BaseBuffer, typing.IO[bytes],
1017 # typing.IO[Any]]"; expected "pandas._typing.IO[Any]"
1018 handle=handle, # type: ignore[arg-type]
1019 # error: Argument "created_handles" to "IOHandles" has incompatible type
1020 # "List[BaseBuffer]"; expected "List[Union[IO[bytes], IO[str]]]"
1021 created_handles=handles, # type: ignore[arg-type]
1022 is_wrapped=is_wrapped,
1023 compression=ioargs.compression,
1024 )
1027class _BufferedWriter(BytesIO, ABC):
1028 """
1029 Some objects do not support multiple .write() calls (TarFile and ZipFile).
1030 This wrapper writes to the underlying buffer on close.
1031 """
1033 buffer = BytesIO()
1035 @abstractmethod
1036 def write_to_buffer(self) -> None: ...
1038 def close(self) -> None:
1039 if self.closed:
1040 # already closed
1041 return
1042 if self.getbuffer().nbytes:
1043 # write to buffer
1044 self.seek(0)
1045 with self.buffer:
1046 self.write_to_buffer()
1047 else:
1048 self.buffer.close()
1049 super().close()
1052class _BytesTarFile(_BufferedWriter):
1053 def __init__(
1054 self,
1055 name: str | None = None,
1056 mode: Literal["r", "a", "w", "x"] = "r",
1057 fileobj: ReadBuffer[bytes] | WriteBuffer[bytes] | None = None,
1058 archive_name: str | None = None,
1059 **kwargs: Any,
1060 ) -> None:
1061 super().__init__()
1062 self.archive_name = archive_name
1063 self.name = name
1064 # error: No overload variant of "open" of "TarFile" matches argument
1065 # types "str | None", "str", "ReadBuffer[bytes] | WriteBuffer[bytes] | None",
1066 # "dict[str, Any]"
1067 # error: Incompatible types in assignment (expression has type "TarFile",
1068 # base class "_BufferedWriter" defined the type as "BytesIO")
1069 self.buffer: tarfile.TarFile = tarfile.TarFile.open( # type: ignore[call-overload, assignment]
1070 name=name,
1071 mode=self.extend_mode(mode),
1072 fileobj=fileobj,
1073 **kwargs,
1074 )
1076 def extend_mode(self, mode: str) -> str:
1077 mode = mode.replace("b", "")
1078 if mode != "w":
1079 return mode
1080 if self.name is not None:
1081 suffix = Path(self.name).suffix
1082 if suffix in (".gz", ".xz", ".bz2"):
1083 mode = f"{mode}:{suffix[1:]}"
1084 return mode
1086 def infer_filename(self) -> str | None:
1087 """
1088 If an explicit archive_name is not given, we still want the file inside the zip
1089 file not to be named something.tar, because that causes confusion (GH39465).
1090 """
1091 if self.name is None:
1092 return None
1094 filename = Path(self.name)
1095 if filename.suffix == ".tar":
1096 return filename.with_suffix("").name
1097 elif filename.suffix in (".tar.gz", ".tar.bz2", ".tar.xz"):
1098 return filename.with_suffix("").with_suffix("").name
1099 return filename.name
1101 def write_to_buffer(self) -> None:
1102 # TarFile needs a non-empty string
1103 archive_name = self.archive_name or self.infer_filename() or "tar"
1104 tarinfo = tarfile.TarInfo(name=archive_name)
1105 tarinfo.size = len(self.getvalue())
1106 self.buffer.addfile(tarinfo, self)
1109class _BytesZipFile(_BufferedWriter):
1110 def __init__(
1111 self,
1112 file: FilePath | ReadBuffer[bytes] | WriteBuffer[bytes],
1113 mode: str,
1114 archive_name: str | None = None,
1115 **kwargs: Any,
1116 ) -> None:
1117 super().__init__()
1118 mode = mode.replace("b", "")
1119 self.archive_name = archive_name
1121 kwargs.setdefault("compression", zipfile.ZIP_DEFLATED)
1122 # error: No overload variant of "ZipFile" matches argument types
1123 # "str | PathLike[str] | ReadBuffer[bytes] | WriteBuffer[bytes]",
1124 # "str", "dict[str, Any]"
1125 # error: Incompatible types in assignment (expression has type "ZipFile",
1126 # base class "_BufferedWriter" defined the type as "BytesIO")
1127 self.buffer: zipfile.ZipFile = zipfile.ZipFile( # type: ignore[call-overload, assignment]
1128 file, mode, **kwargs
1129 )
1131 def infer_filename(self) -> str | None:
1132 """
1133 If an explicit archive_name is not given, we still want the file inside the zip
1134 file not to be named something.zip, because that causes confusion (GH39465).
1135 """
1136 if isinstance(self.buffer.filename, (os.PathLike, str)):
1137 filename = Path(self.buffer.filename)
1138 if filename.suffix == ".zip":
1139 return filename.with_suffix("").name
1140 return filename.name
1141 return None
1143 def write_to_buffer(self) -> None:
1144 # ZipFile needs a non-empty string
1145 archive_name = self.archive_name or self.infer_filename() or "zip"
1146 self.buffer.writestr(archive_name, self.getvalue())
1149class _IOWrapper:
1150 # TextIOWrapper is overly strict: it request that the buffer has seekable, readable,
1151 # and writable. If we have a read-only buffer, we shouldn't need writable and vice
1152 # versa. Some buffers, are seek/read/writ-able but they do not have the "-able"
1153 # methods, e.g., tempfile.SpooledTemporaryFile.
1154 # If a buffer does not have the above "-able" methods, we simply assume they are
1155 # seek/read/writ-able.
1156 def __init__(self, buffer: BaseBuffer) -> None:
1157 self.buffer = buffer
1159 def __getattr__(self, name: str) -> Any:
1160 return getattr(self.buffer, name)
1162 def readable(self) -> bool:
1163 if hasattr(self.buffer, "readable"):
1164 return self.buffer.readable()
1165 return True
1167 def seekable(self) -> bool:
1168 if hasattr(self.buffer, "seekable"):
1169 return self.buffer.seekable()
1170 return True
1172 def writable(self) -> bool:
1173 if hasattr(self.buffer, "writable"):
1174 return self.buffer.writable()
1175 return True
1178class _BytesIOWrapper:
1179 # Wrapper that wraps a StringIO buffer and reads bytes from it
1180 # Created for compat with pyarrow read_csv
1181 def __init__(
1182 self,
1183 buffer: StringIO | TextIOBase,
1184 encoding: str = "utf-8",
1185 errors: str = "strict",
1186 ) -> None:
1187 self.buffer = buffer
1188 self.encoding = encoding
1189 self.errors = errors
1190 # Because a character can be represented by more than 1 byte,
1191 # it is possible that reading will produce more bytes than n
1192 # We store the extra bytes in this overflow variable, and append the
1193 # overflow to the front of the bytestring the next time reading is performed
1194 self.overflow = b""
1196 def __getattr__(self, attr: str) -> Any:
1197 return getattr(self.buffer, attr)
1199 def read(self, n: int | None = -1) -> bytes:
1200 assert self.buffer is not None
1201 bytestring = self.buffer.read(n).encode(self.encoding, self.errors)
1202 # When n=-1/n greater than remaining bytes: Read entire file/rest of file
1203 combined_bytestring = self.overflow + bytestring
1204 if n is None or n < 0 or n >= len(combined_bytestring):
1205 self.overflow = b""
1206 return combined_bytestring
1207 else:
1208 to_return = combined_bytestring[:n]
1209 self.overflow = combined_bytestring[n:]
1210 return to_return
1213def _maybe_memory_map(
1214 handle: str | BaseBuffer, memory_map: bool
1215) -> tuple[str | BaseBuffer, bool, list[BaseBuffer]]:
1216 """Try to memory map file/buffer."""
1217 handles: list[BaseBuffer] = []
1218 memory_map &= hasattr(handle, "fileno") or isinstance(handle, str)
1219 if not memory_map:
1220 return handle, memory_map, handles
1222 # mmap used by only read_csv
1223 handle = cast("ReadCsvBuffer[bytes]", handle)
1225 # need to open the file first
1226 if isinstance(handle, str):
1227 handle = open(handle, "rb")
1228 handles.append(handle)
1230 try:
1231 # open mmap and adds *-able
1232 # error: Argument 1 to "_IOWrapper" has incompatible type "mmap";
1233 # expected "BaseBuffer"
1234 wrapped = _IOWrapper(
1235 mmap.mmap(
1236 handle.fileno(),
1237 0,
1238 access=mmap.ACCESS_READ, # type: ignore[arg-type]
1239 )
1240 )
1241 except UnsupportedOperation as err:
1242 # GH#45630 in-memory buffers like BytesIO/StringIO have a fileno
1243 # method but raise UnsupportedOperation when it is called
1244 raise ValueError(
1245 "memory_map=True is only supported when reading from a file path "
1246 "or a file-like object backed by a real file descriptor; "
1247 "in-memory buffers (e.g. BytesIO, StringIO) are not supported."
1248 ) from err
1249 finally:
1250 for handle in reversed(handles):
1251 # error: "BaseBuffer" has no attribute "close"
1252 handle.close() # type: ignore[attr-defined]
1254 return wrapped, memory_map, [wrapped]
1257def file_exists(filepath_or_buffer: FilePath | BaseBuffer) -> bool:
1258 """Test whether file exists."""
1259 exists = False
1260 filepath_or_buffer = stringify_path(filepath_or_buffer)
1261 if not isinstance(filepath_or_buffer, str):
1262 return exists
1263 try:
1264 exists = os.path.exists(filepath_or_buffer)
1265 # gh-5874: if the filepath is too long will raise here
1266 except (TypeError, ValueError):
1267 pass
1268 return exists
1271def _is_binary_mode(handle: FilePath | BaseBuffer, mode: str) -> bool:
1272 """Whether the handle is opened in binary mode"""
1273 # specified by user
1274 if "t" in mode or "b" in mode:
1275 return "b" in mode
1277 # exceptions
1278 text_classes = (
1279 # classes that expect string but have 'b' in mode
1280 codecs.StreamWriter,
1281 codecs.StreamReader,
1282 codecs.StreamReaderWriter,
1283 TextIOBase,
1284 )
1285 if issubclass(type(handle), text_classes):
1286 return False
1288 if isinstance(handle, _get_binary_io_classes()):
1289 return True
1291 if hasattr(handle, "mode"):
1292 return "b" in handle.mode
1294 # GH#52252 buffers without a "mode" attribute that are not covered by the
1295 # classes above, e.g. mmap.mmap and botocore's StreamingBody, which is an
1296 # IOBase subclass without being a Raw/BufferedIOBase subclass.
1297 return isinstance(handle, (mmap.mmap, IOBase))
1300@functools.lru_cache
1301def _get_binary_io_classes() -> tuple[type, ...]:
1302 """IO classes that expect bytes"""
1303 binary_classes: tuple[type, ...] = (BufferedIOBase, RawIOBase)
1305 # python-zstandard doesn't use any of the builtin base classes; instead we
1306 # have to use the `zstd.ZstdDecompressionReader` class for isinstance checks.
1307 # Unfortunately `zstd.ZstdDecompressionReader` isn't exposed by python-zstandard
1308 # so we have to get it from a `zstd.ZstdDecompressor` instance.
1309 # See also https://github.com/indygreg/python-zstandard/pull/165.
1310 zstd = import_optional_dependency("zstandard", errors="ignore")
1311 if zstd is not None:
1312 with zstd.ZstdDecompressor().stream_reader(b"") as reader:
1313 binary_classes += (type(reader),)
1315 return binary_classes
1318def is_potential_multi_index(
1319 columns: Sequence[Hashable] | MultiIndex,
1320 index_col: bool | Sequence[int] | None = None,
1321) -> bool:
1322 """
1323 Check whether or not the `columns` parameter
1324 could be converted into a MultiIndex.
1326 Parameters
1327 ----------
1328 columns : array-like
1329 Object which may or may not be convertible into a MultiIndex
1330 index_col : None, bool or list, optional
1331 Column or columns to use as the (possibly hierarchical) index
1333 Returns
1334 -------
1335 bool : Whether or not columns could become a MultiIndex
1336 """
1337 if index_col is None or isinstance(index_col, bool):
1338 index_columns = set()
1339 else:
1340 index_columns = set(index_col)
1342 return bool(
1343 len(columns)
1344 and not isinstance(columns, ABCMultiIndex)
1345 and all(isinstance(c, tuple) for c in columns if c not in index_columns)
1346 )
1349def dedup_names(
1350 names: Sequence[Hashable], is_potential_multiindex: bool
1351) -> Sequence[Hashable]:
1352 """
1353 Rename column names if duplicates exist.
1355 Currently the renaming is done by appending a period and an autonumeric,
1356 but a custom pattern may be supported in the future.
1358 Examples
1359 --------
1360 >>> dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False)
1361 ['x', 'y', 'x.1', 'x.2']
1362 """
1363 names = list(names) # so we can index
1364 counts: defaultdict[Hashable, int] = defaultdict(int)
1366 for i, col in enumerate(names):
1367 cur_count = counts[col]
1369 while cur_count > 0:
1370 counts[col] = cur_count + 1
1372 if is_potential_multiindex:
1373 # for mypy
1374 assert isinstance(col, tuple)
1375 col = (*col[:-1], f"{col[-1]}.{cur_count}")
1376 else:
1377 col = f"{col}.{cur_count}"
1378 cur_count = counts[col]
1380 names[i] = col
1381 counts[col] = cur_count + 1
1383 return names
1386def mangle_dupe_names(
1387 names: Sequence[Hashable],
1388 unnamed_col_indices: Sequence[int] = (),
1389 dtype: DtypeArg | None = None,
1390) -> list[Hashable]:
1391 """
1392 De-duplicate column names, matching ``pandas._libs.parsers.TextReader``.
1394 This is the header-mangling used by the read_csv engines, and it differs
1395 from :func:`dedup_names` in three ways (see GH#50371):
1397 - a mangled name that would collide with a name already present in
1398 ``names`` is skipped, so ``["x", "x", "x.1"]`` becomes
1399 ``["x", "x.2", "x.1"]`` rather than ``["x", "x.1", "x.1.1"]``;
1400 - columns whose position is listed in ``unnamed_col_indices`` are mangled
1401 only after the named columns, so the named columns keep their name;
1402 - a dict ``dtype`` is updated in place so that a renamed column keeps the
1403 dtype requested for its original name.
1405 Examples
1406 --------
1407 >>> mangle_dupe_names(["x", "x", "x.1"])
1408 ['x', 'x.2', 'x.1']
1409 """
1410 names = list(names) # so we can index
1411 unnamed = set(unnamed_col_indices)
1412 # Ensure that regular columns are used before unnamed ones
1413 # to keep given names and mangle unnamed columns
1414 col_loop_order = [i for i in range(len(names)) if i not in unnamed] + list(
1415 unnamed_col_indices
1416 )
1417 counts: dict[Hashable, int] = {}
1419 for i in col_loop_order:
1420 col = old_col = names[i]
1421 cur_count = counts.get(col, 0)
1423 if cur_count > 0:
1424 while cur_count > 0:
1425 counts[old_col] = cur_count + 1
1426 col = f"{old_col}.{cur_count}"
1427 if col in names:
1428 cur_count += 1
1429 else:
1430 cur_count = counts.get(col, 0)
1432 if (
1433 isinstance(dtype, dict)
1434 and dtype.get(old_col) is not None
1435 and dtype.get(col) is None
1436 ):
1437 dtype[col] = dtype[old_col]
1439 names[i] = col
1440 counts[col] = cur_count + 1
1442 return names