Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/common.py: 23%
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 RawIOBase,
23 StringIO,
24 TextIOBase,
25 TextIOWrapper,
26)
27import mmap
28import os
29from pathlib import Path
30import re
31import tarfile
32from typing import (
33 IO,
34 TYPE_CHECKING,
35 Any,
36 AnyStr,
37 DefaultDict,
38 Generic,
39 Literal,
40 TypeVar,
41 cast,
42 overload,
43)
44from urllib.parse import (
45 urljoin,
46 urlparse as parse_url,
47 uses_netloc,
48 uses_params,
49 uses_relative,
50)
51import warnings
52import zipfile
54from pandas._typing import (
55 BaseBuffer,
56 ReadCsvBuffer,
57)
58from pandas.compat._optional import import_optional_dependency
59from pandas.util._exceptions import find_stack_level
61from pandas.core.dtypes.common import (
62 is_bool,
63 is_file_like,
64 is_integer,
65 is_list_like,
66)
67from pandas.core.dtypes.generic import ABCMultiIndex
69_VALID_URLS = set(uses_relative + uses_netloc + uses_params)
70_VALID_URLS.discard("")
71_FSSPEC_URL_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9+\-+.]*(::[A-Za-z0-9+\-+.]+)*://")
73BaseBufferT = TypeVar("BaseBufferT", bound=BaseBuffer)
76if TYPE_CHECKING:
77 from types import TracebackType
79 from pandas._typing import (
80 CompressionDict,
81 CompressionOptions,
82 FilePath,
83 ReadBuffer,
84 StorageOptions,
85 WriteBuffer,
86 )
88 from pandas import MultiIndex
91@dataclasses.dataclass
92class IOArgs:
93 """
94 Return value of io/common.py:_get_filepath_or_buffer.
95 """
97 filepath_or_buffer: str | BaseBuffer
98 encoding: str
99 mode: str
100 compression: CompressionDict
101 should_close: bool = False
102 close_handles: list[Any] = dataclasses.field(default_factory=list)
105@dataclasses.dataclass
106class IOHandles(Generic[AnyStr]):
107 """
108 Return value of io/common.py:get_handle
110 Can be used as a context manager.
112 This is used to easily close created buffers and to handle corner cases when
113 TextIOWrapper is inserted.
115 handle: The file handle to be used.
116 created_handles: All file handles that are created by get_handle
117 is_wrapped: Whether a TextIOWrapper needs to be detached.
118 """
120 # handle might not implement the IO-interface
121 handle: IO[AnyStr]
122 compression: CompressionDict
123 created_handles: list[IO[bytes] | IO[str]] = dataclasses.field(default_factory=list)
124 is_wrapped: bool = False
126 def close(self) -> None:
127 """
128 Close all created buffers.
130 Note: If a TextIOWrapper was inserted, it is flushed and detached to
131 avoid closing the potentially user-created buffer.
132 """
133 if self.is_wrapped:
134 assert isinstance(self.handle, TextIOWrapper)
135 self.handle.flush()
136 self.handle.detach()
137 self.created_handles.remove(self.handle)
138 for handle in self.created_handles:
139 handle.close()
140 self.created_handles = []
141 self.is_wrapped = False
143 def __enter__(self) -> IOHandles[AnyStr]:
144 return self
146 def __exit__(
147 self,
148 exc_type: type[BaseException] | None,
149 exc_value: BaseException | None,
150 traceback: TracebackType | None,
151 ) -> None:
152 self.close()
155def is_url(url: object) -> bool:
156 """
157 Check to see if a URL has a valid protocol.
159 Parameters
160 ----------
161 url : str or unicode
163 Returns
164 -------
165 isurl : bool
166 If `url` has a valid protocol return True otherwise False.
167 """
168 if not isinstance(url, str):
169 return False
170 return parse_url(url).scheme in _VALID_URLS
173@overload
174def _expand_user(filepath_or_buffer: str) -> str: ...
177@overload
178def _expand_user(filepath_or_buffer: BaseBufferT) -> BaseBufferT: ...
181def _expand_user(filepath_or_buffer: str | BaseBufferT) -> str | BaseBufferT:
182 """
183 Return the argument with an initial component of ~ or ~user
184 replaced by that user's home directory.
186 Parameters
187 ----------
188 filepath_or_buffer : object to be converted if possible
190 Returns
191 -------
192 expanded_filepath_or_buffer : an expanded filepath or the
193 input if not expandable
194 """
195 if isinstance(filepath_or_buffer, str):
196 return os.path.expanduser(filepath_or_buffer)
197 return filepath_or_buffer
200def validate_header_arg(header: object) -> None:
201 if header is None:
202 return
203 if is_integer(header):
204 header = cast(int, header)
205 if header < 0:
206 # GH 27779
207 raise ValueError(
208 "Passing negative integer to header is invalid. "
209 "For no header, use header=None instead"
210 )
211 return
212 if is_list_like(header, allow_sets=False):
213 header = cast(Sequence, header)
214 if not all(map(is_integer, header)):
215 raise ValueError("header must be integer or list of integers")
216 if any(i < 0 for i in header):
217 raise ValueError("cannot specify multi-index header with negative integers")
218 return
219 if is_bool(header):
220 raise TypeError(
221 "Passing a bool to header is invalid. Use header=None for no header or "
222 "header=int or list-like of ints to specify "
223 "the row(s) making up the column names"
224 )
225 # GH 16338
226 raise ValueError("header must be integer or list of integers")
229@overload
230def stringify_path(
231 filepath_or_buffer: FilePath, convert_file_like: bool = ...
232) -> str: ...
235@overload
236def stringify_path(
237 filepath_or_buffer: BaseBufferT, convert_file_like: bool = ...
238) -> BaseBufferT: ...
241def stringify_path(
242 filepath_or_buffer: FilePath | BaseBufferT,
243 convert_file_like: bool = False,
244) -> str | BaseBufferT:
245 """
246 Attempt to convert a path-like object to a string.
248 Parameters
249 ----------
250 filepath_or_buffer : object to be converted
252 Returns
253 -------
254 str_filepath_or_buffer : maybe a string version of the object
256 Notes
257 -----
258 Objects supporting the fspath protocol are coerced
259 according to its __fspath__ method.
261 Any other object is passed through unchanged, which includes bytes,
262 strings, buffers, or anything else that's not even path-like.
263 """
264 if not convert_file_like and is_file_like(filepath_or_buffer):
265 # GH 38125: some fsspec objects implement os.PathLike but have already opened a
266 # file. This prevents opening the file a second time. infer_compression calls
267 # this function with convert_file_like=True to infer the compression.
268 return cast(BaseBufferT, filepath_or_buffer)
270 if isinstance(filepath_or_buffer, os.PathLike):
271 filepath_or_buffer = filepath_or_buffer.__fspath__()
272 return _expand_user(filepath_or_buffer)
275def urlopen(*args: Any, **kwargs: Any) -> Any:
276 """
277 Lazy-import wrapper for stdlib urlopen, as that imports a big chunk of
278 the stdlib.
279 """
280 import urllib.request
282 return urllib.request.urlopen(*args, **kwargs) # noqa: TID251
285def is_fsspec_url(url: FilePath | BaseBuffer) -> bool:
286 """
287 Returns true if the given URL looks like
288 something fsspec can handle
289 """
290 return (
291 isinstance(url, str)
292 and bool(_FSSPEC_URL_PATTERN.match(url))
293 and not url.startswith(("http://", "https://"))
294 )
297def _get_filepath_or_buffer(
298 filepath_or_buffer: FilePath | BaseBuffer,
299 encoding: str = "utf-8",
300 compression: CompressionOptions | None = None,
301 mode: str = "r",
302 storage_options: StorageOptions | None = None,
303) -> IOArgs:
304 """
305 If the filepath_or_buffer is a url, translate and return the buffer.
306 Otherwise passthrough.
308 Parameters
309 ----------
310 filepath_or_buffer : a url, filepath (str or pathlib.Path),
311 or buffer
313 compression : str or dict, default 'infer'
314 For on-the-fly compression of the output data. If 'infer' and
315 'filepath_or_buffer' is path-like, then detect compression from the
316 following extensions: '.gz',
317 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
318 (otherwise no compression).
319 Set to ``None`` for no compression.
320 Can also be a dict with key ``'method'`` set
321 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
322 and other key-value pairs are forwarded to
323 ``zipfile.ZipFile``, ``gzip.GzipFile``,
324 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
325 ``tarfile.TarFile``, respectively.
326 As an example, the following could be passed for faster compression and to
327 create a reproducible gzip archive:
328 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
330 encoding : the encoding to use to decode bytes, default is 'utf-8'
331 mode : str, optional
333 storage_options : dict, optional
334 Extra options that make sense for a particular storage connection, e.g.
335 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs
336 are forwarded to ``urllib.request.Request`` as header options. For other
337 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are
338 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more
339 details, and for more examples on storage options refer `here
340 <https://pandas.pydata.org/docs/user_guide/io.html?
341 highlight=storage_options#reading-writing-remote-files>`_.
344 Returns the dataclass IOArgs.
345 """
346 filepath_or_buffer = stringify_path(filepath_or_buffer)
348 # handle compression dict
349 compression_method, compression = get_compression_method(compression)
350 compression_method = infer_compression(filepath_or_buffer, compression_method)
352 # GH21227 internal compression is not used for non-binary handles.
353 if compression_method and hasattr(filepath_or_buffer, "write") and "b" not in mode:
354 warnings.warn(
355 "compression has no effect when passing a non-binary object as input.",
356 RuntimeWarning,
357 stacklevel=find_stack_level(),
358 )
359 compression_method = None
361 compression = dict(compression, method=compression_method)
363 # bz2 and xz do not write the byte order mark for utf-16 and utf-32
364 # print a warning when writing such files
365 if (
366 "w" in mode
367 and compression_method in ["bz2", "xz"]
368 and encoding in ["utf-16", "utf-32"]
369 ):
370 warnings.warn(
371 f"{compression} will not write the byte order mark for {encoding}",
372 UnicodeWarning,
373 stacklevel=find_stack_level(),
374 )
376 if "a" in mode and compression_method in ["zip", "tar"]:
377 # GH56778
378 warnings.warn(
379 "zip and tar do not support mode 'a' properly. "
380 "This combination will result in multiple files with same name "
381 "being added to the archive.",
382 RuntimeWarning,
383 stacklevel=find_stack_level(),
384 )
386 # Use binary mode when converting path-like objects to file-like objects (fsspec)
387 # except when text mode is explicitly requested. The original mode is returned if
388 # fsspec is not used.
389 fsspec_mode = mode
390 if "t" not in fsspec_mode and "b" not in fsspec_mode:
391 fsspec_mode += "b"
393 if isinstance(filepath_or_buffer, str) and is_url(filepath_or_buffer):
394 # TODO: fsspec can also handle HTTP via requests, but leaving this
395 # unchanged. using fsspec appears to break the ability to infer if the
396 # server responded with gzipped data
397 storage_options = storage_options or {}
399 # waiting until now for importing to match intended lazy logic of
400 # urlopen function defined elsewhere in this module
401 import urllib.request
403 # assuming storage_options is to be interpreted as headers
404 req_info = urllib.request.Request(filepath_or_buffer, headers=storage_options)
405 with urlopen(req_info) as req:
406 content_encoding = req.headers.get("Content-Encoding", None)
407 if content_encoding == "gzip":
408 # Override compression based on Content-Encoding header
409 compression = {"method": "gzip"}
410 reader = BytesIO(req.read())
411 return IOArgs(
412 filepath_or_buffer=reader,
413 encoding=encoding,
414 compression=compression,
415 should_close=True,
416 mode=fsspec_mode,
417 )
419 if is_fsspec_url(filepath_or_buffer):
420 assert isinstance(
421 filepath_or_buffer, str
422 ) # just to appease mypy for this branch
423 # two special-case s3-like protocols; these have special meaning in Hadoop,
424 # but are equivalent to just "s3" from fsspec's point of view
425 # cc #11071
426 if filepath_or_buffer.startswith("s3a://"):
427 filepath_or_buffer = filepath_or_buffer.replace("s3a://", "s3://")
428 if filepath_or_buffer.startswith("s3n://"):
429 filepath_or_buffer = filepath_or_buffer.replace("s3n://", "s3://")
430 fsspec = import_optional_dependency("fsspec")
432 # If botocore is installed we fallback to reading with anon=True
433 # to allow reads from public buckets
434 err_types_to_retry_with_anon: list[Any] = []
435 try:
436 import_optional_dependency("botocore")
437 from botocore.exceptions import (
438 ClientError,
439 NoCredentialsError,
440 )
442 err_types_to_retry_with_anon = [
443 ClientError,
444 NoCredentialsError,
445 PermissionError,
446 ]
447 except ImportError:
448 pass
450 try:
451 open_file = fsspec.open(
452 filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
453 )
454 file_obj = open_file.open()
455 # GH 34626 Reads from Public Buckets without Credentials needs anon=True
456 except tuple(err_types_to_retry_with_anon):
457 if storage_options is None:
458 storage_options = {"anon": True}
459 else:
460 # don't mutate user input.
461 storage_options = dict(storage_options)
462 storage_options["anon"] = True
463 open_file = fsspec.open(
464 filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
465 )
466 file_obj = open_file.open()
468 return IOArgs(
469 filepath_or_buffer=file_obj,
470 encoding=encoding,
471 compression=compression,
472 close_handles=[open_file],
473 should_close=True,
474 mode=fsspec_mode,
475 )
476 elif storage_options:
477 raise ValueError(
478 "storage_options passed with file object or non-fsspec file path"
479 )
481 if isinstance(filepath_or_buffer, (str, bytes, mmap.mmap)):
482 return IOArgs(
483 filepath_or_buffer=_expand_user(filepath_or_buffer),
484 encoding=encoding,
485 compression=compression,
486 should_close=False,
487 mode=mode,
488 )
490 # is_file_like requires (read | write) & __iter__ but __iter__ is only
491 # needed for read_csv(engine=python)
492 if not (
493 hasattr(filepath_or_buffer, "read") or hasattr(filepath_or_buffer, "write")
494 ):
495 msg = f"Invalid file path or buffer object type: {type(filepath_or_buffer)}"
496 raise ValueError(msg)
498 return IOArgs(
499 filepath_or_buffer=filepath_or_buffer,
500 encoding=encoding,
501 compression=compression,
502 should_close=False,
503 mode=mode,
504 )
507def file_path_to_url(path: str) -> str:
508 """
509 converts an absolute native path to a FILE URL.
511 Parameters
512 ----------
513 path : a path in native format
515 Returns
516 -------
517 a valid FILE URL
518 """
519 # lazify expensive import (~30ms)
520 from urllib.request import pathname2url
522 return urljoin("file:", pathname2url(path))
525extension_to_compression = {
526 ".tar": "tar",
527 ".tar.gz": "tar",
528 ".tar.bz2": "tar",
529 ".tar.xz": "tar",
530 ".gz": "gzip",
531 ".bz2": "bz2",
532 ".zip": "zip",
533 ".xz": "xz",
534 ".zst": "zstd",
535}
536_supported_compressions = set(extension_to_compression.values())
539def get_compression_method(
540 compression: CompressionOptions,
541) -> tuple[str | None, CompressionDict]:
542 """
543 Simplifies a compression argument to a compression method string and
544 a mapping containing additional arguments.
546 Parameters
547 ----------
548 compression : str or mapping
549 If string, specifies the compression method. If mapping, value at key
550 'method' specifies compression method.
552 Returns
553 -------
554 tuple of ({compression method}, Optional[str]
555 {compression arguments}, Dict[str, Any])
557 Raises
558 ------
559 ValueError on mapping missing 'method' key
560 """
561 compression_method: str | None
562 if isinstance(compression, Mapping):
563 compression_args = dict(compression)
564 try:
565 compression_method = compression_args.pop("method")
566 except KeyError as err:
567 raise ValueError("If mapping, compression must have key 'method'") from err
568 else:
569 compression_args = {}
570 compression_method = compression
571 return compression_method, compression_args
574def infer_compression(
575 filepath_or_buffer: FilePath | BaseBuffer, compression: str | None
576) -> str | None:
577 """
578 Get the compression method for filepath_or_buffer. If compression='infer',
579 the inferred compression method is returned. Otherwise, the input
580 compression method is returned unchanged, unless it's invalid, in which
581 case an error is raised.
583 Parameters
584 ----------
585 filepath_or_buffer : str or file handle
586 File path or object.
588 compression : str or dict, default 'infer'
589 For on-the-fly compression of the output data. If 'infer' and
590 'filepath_or_buffer' is path-like, then detect compression from the
591 following extensions: '.gz',
592 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
593 (otherwise no compression).
594 Set to ``None`` for no compression.
595 Can also be a dict with key ``'method'`` set
596 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
597 and other key-value pairs are forwarded to
598 ``zipfile.ZipFile``, ``gzip.GzipFile``,
599 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
600 ``tarfile.TarFile``, respectively.
601 As an example, the following could be passed for faster compression and to
602 create a reproducible gzip archive:
603 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
605 Returns
606 -------
607 string or None
609 Raises
610 ------
611 ValueError on invalid compression specified.
612 """
613 if compression is None:
614 return None
616 # Infer compression
617 if compression == "infer":
618 # Convert all path types (e.g. pathlib.Path) to strings
619 if isinstance(filepath_or_buffer, str) and "::" in filepath_or_buffer:
620 # chained URLs contain ::
621 filepath_or_buffer = filepath_or_buffer.split("::")[0]
622 filepath_or_buffer = stringify_path(filepath_or_buffer, convert_file_like=True)
623 if not isinstance(filepath_or_buffer, str):
624 # Cannot infer compression of a buffer, assume no compression
625 return None
627 # Infer compression from the filename/URL extension
628 for extension, compression in extension_to_compression.items():
629 if filepath_or_buffer.lower().endswith(extension):
630 return compression
631 return None
633 # Compression has been specified. Check that it's valid
634 if compression in _supported_compressions:
635 return compression
637 valid = ["infer", None, *sorted(_supported_compressions)]
638 msg = (
639 f"Unrecognized compression type: {compression}\n"
640 f"Valid compression types are {valid}"
641 )
642 raise ValueError(msg)
645def check_parent_directory(path: Path | str) -> None:
646 """
647 Check if parent directory of a file exists, raise OSError if it does not
649 Parameters
650 ----------
651 path: Path or str
652 Path to check parent directory of
653 """
654 parent = Path(path).parent
655 if not parent.is_dir():
656 raise OSError(rf"Cannot save file into a non-existent directory: '{parent}'")
659@overload
660def get_handle(
661 path_or_buf: FilePath | BaseBuffer,
662 mode: str,
663 *,
664 encoding: str | None = ...,
665 compression: CompressionOptions = ...,
666 memory_map: bool = ...,
667 is_text: Literal[False],
668 errors: str | None = ...,
669 storage_options: StorageOptions = ...,
670) -> IOHandles[bytes]: ...
673@overload
674def get_handle(
675 path_or_buf: FilePath | BaseBuffer,
676 mode: str,
677 *,
678 encoding: str | None = ...,
679 compression: CompressionOptions = ...,
680 memory_map: bool = ...,
681 is_text: Literal[True] = ...,
682 errors: str | None = ...,
683 storage_options: StorageOptions = ...,
684) -> IOHandles[str]: ...
687@overload
688def get_handle(
689 path_or_buf: FilePath | BaseBuffer,
690 mode: str,
691 *,
692 encoding: str | None = ...,
693 compression: CompressionOptions = ...,
694 memory_map: bool = ...,
695 is_text: bool = ...,
696 errors: str | None = ...,
697 storage_options: StorageOptions = ...,
698) -> IOHandles[str] | IOHandles[bytes]: ...
701def get_handle(
702 path_or_buf: FilePath | BaseBuffer,
703 mode: str,
704 *,
705 encoding: str | None = None,
706 compression: CompressionOptions | None = None,
707 memory_map: bool = False,
708 is_text: bool = True,
709 errors: str | None = None,
710 storage_options: StorageOptions | None = None,
711) -> IOHandles[str] | IOHandles[bytes]:
712 """
713 Get file handle for given path/buffer and mode.
715 Parameters
716 ----------
717 path_or_buf : str or file handle
718 File path or object.
719 mode : str
720 Mode to open path_or_buf with.
721 encoding : str or None
722 Encoding to use.
723 compression : str or dict, default 'infer'
724 For on-the-fly compression of the output data. If 'infer' and 'path_or_buf'
725 is path-like, then detect compression from the following extensions: '.gz',
726 '.bz2', '.zip', '.xz', '.zst', '.tar', '.tar.gz', '.tar.xz' or '.tar.bz2'
727 (otherwise no compression).
728 Set to ``None`` for no compression.
729 Can also be a dict with key ``'method'`` set
730 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, ``'tar'``}
731 and other key-value pairs are forwarded to
732 ``zipfile.ZipFile``, ``gzip.GzipFile``,
733 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or
734 ``tarfile.TarFile``, respectively.
735 As an example, the following could be passed for faster compression and to
736 create a reproducible gzip archive:
737 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
739 May be a dict with key 'method' as compression mode
740 and other keys as compression options if compression
741 mode is 'zip'.
743 Passing compression options as keys in dict is
744 supported for compression modes 'gzip', 'bz2', 'zstd' and 'zip'.
746 memory_map : bool, default False
747 See parsers._parser_params for more information. Only used by read_csv.
748 is_text : bool, default True
749 Whether the type of the content passed to the file/buffer is string or
750 bytes. This is not the same as `"b" not in mode`. If a string content is
751 passed to a binary file/buffer, a wrapper is inserted.
752 errors : str, default 'strict'
753 Specifies how encoding and decoding errors are to be handled.
754 See the errors argument for :func:`open` for a full list
755 of options.
756 storage_options: StorageOptions = None
757 Passed to _get_filepath_or_buffer
759 Returns the dataclass IOHandles
760 """
761 # Windows does not default to utf-8. Set to utf-8 for a consistent behavior
762 encoding = encoding or "utf-8"
764 errors = errors or "strict"
766 # read_csv does not know whether the buffer is opened in binary/text mode
767 if _is_binary_mode(path_or_buf, mode) and "b" not in mode:
768 mode += "b"
770 # validate encoding and errors
771 codecs.lookup(encoding)
772 if isinstance(errors, str):
773 codecs.lookup_error(errors)
775 # open URLs
776 ioargs = _get_filepath_or_buffer(
777 path_or_buf,
778 encoding=encoding,
779 compression=compression,
780 mode=mode,
781 storage_options=storage_options,
782 )
784 handle = ioargs.filepath_or_buffer
785 handles: list[BaseBuffer]
787 # memory mapping needs to be the first step
788 # only used for read_csv
789 handle, memory_map, handles = _maybe_memory_map(handle, memory_map)
791 is_path = isinstance(handle, str)
792 compression_args = dict(ioargs.compression)
793 compression = compression_args.pop("method")
795 # Only for write methods
796 if "r" not in mode and is_path:
797 check_parent_directory(str(handle))
799 if compression:
800 if compression != "zstd":
801 # compression libraries do not like an explicit text-mode
802 ioargs.mode = ioargs.mode.replace("t", "")
803 elif compression == "zstd" and "b" not in ioargs.mode:
804 # python-zstandard defaults to text mode, but we always expect
805 # compression libraries to use binary mode.
806 ioargs.mode += "b"
808 # GZ Compression
809 if compression == "gzip":
810 if isinstance(handle, str):
811 # error: Incompatible types in assignment (expression has type
812 # "GzipFile", variable has type "Union[str, BaseBuffer]")
813 handle = gzip.GzipFile( # type: ignore[assignment]
814 filename=handle,
815 mode=ioargs.mode,
816 **compression_args,
817 )
818 else:
819 handle = gzip.GzipFile(
820 # No overload variant of "GzipFile" matches argument types
821 # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
822 fileobj=handle, # type: ignore[call-overload]
823 mode=ioargs.mode,
824 **compression_args,
825 )
827 # BZ Compression
828 elif compression == "bz2":
829 import bz2
831 # Overload of "BZ2File" to handle pickle protocol 5
832 # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
833 handle = bz2.BZ2File( # type: ignore[call-overload]
834 handle,
835 mode=ioargs.mode,
836 **compression_args,
837 )
839 # ZIP Compression
840 elif compression == "zip":
841 # error: Argument 1 to "_BytesZipFile" has incompatible type
842 # "Union[str, BaseBuffer]"; expected "Union[Union[str, PathLike[str]],
843 # ReadBuffer[bytes], WriteBuffer[bytes]]"
844 handle = _BytesZipFile(
845 handle, # type: ignore[arg-type]
846 ioargs.mode,
847 **compression_args,
848 )
849 if handle.buffer.mode == "r":
850 handles.append(handle)
851 zip_names = handle.buffer.namelist()
852 if len(zip_names) == 1:
853 handle = handle.buffer.open(zip_names.pop())
854 elif not zip_names:
855 raise ValueError(f"Zero files found in ZIP file {path_or_buf}")
856 else:
857 raise ValueError(
858 "Multiple files found in ZIP file. "
859 f"Only one file per ZIP: {zip_names}"
860 )
862 # TAR Encoding
863 elif compression == "tar":
864 compression_args.setdefault("mode", ioargs.mode)
865 if isinstance(handle, str):
866 handle = _BytesTarFile(name=handle, **compression_args)
867 else:
868 # error: Argument "fileobj" to "_BytesTarFile" has incompatible
869 # type "BaseBuffer"; expected "Union[ReadBuffer[bytes],
870 # WriteBuffer[bytes], None]"
871 handle = _BytesTarFile(
872 fileobj=handle, # type: ignore[arg-type]
873 **compression_args,
874 )
875 assert isinstance(handle, _BytesTarFile)
876 if "r" in handle.buffer.mode:
877 handles.append(handle)
878 files = handle.buffer.getnames()
879 if len(files) == 1:
880 file = handle.buffer.extractfile(files[0])
881 assert file is not None
882 handle = file
883 elif not files:
884 raise ValueError(f"Zero files found in TAR archive {path_or_buf}")
885 else:
886 raise ValueError(
887 "Multiple files found in TAR archive. "
888 f"Only one file per TAR archive: {files}"
889 )
891 # XZ Compression
892 elif compression == "xz":
893 # error: Argument 1 to "LZMAFile" has incompatible type "Union[str,
894 # BaseBuffer]"; expected "Optional[Union[Union[str, bytes, PathLike[str],
895 # PathLike[bytes]], IO[bytes]], None]"
896 import lzma
898 handle = lzma.LZMAFile(
899 handle, # type: ignore[arg-type]
900 ioargs.mode,
901 **compression_args,
902 )
904 # Zstd Compression
905 elif compression == "zstd":
906 zstd = import_optional_dependency("zstandard")
907 if "r" in ioargs.mode:
908 open_args = {"dctx": zstd.ZstdDecompressor(**compression_args)}
909 else:
910 open_args = {"cctx": zstd.ZstdCompressor(**compression_args)}
911 handle = zstd.open(
912 handle,
913 mode=ioargs.mode,
914 **open_args,
915 )
917 # Unrecognized Compression
918 else:
919 msg = f"Unrecognized compression type: {compression}"
920 raise ValueError(msg)
922 assert not isinstance(handle, str)
923 handles.append(handle)
925 elif isinstance(handle, str):
926 # Check whether the filename is to be opened in binary mode.
927 # Binary mode does not support 'encoding' and 'newline'.
928 if ioargs.encoding and "b" not in ioargs.mode:
929 # Encoding
930 handle = open(
931 handle,
932 ioargs.mode,
933 encoding=ioargs.encoding,
934 errors=errors,
935 newline="",
936 )
937 else:
938 # Binary mode
939 handle = open(handle, ioargs.mode)
940 handles.append(handle)
942 # Convert BytesIO or file objects passed with an encoding
943 is_wrapped = False
944 if not is_text and ioargs.mode == "rb" and isinstance(handle, TextIOBase):
945 # not added to handles as it does not open/buffer resources
946 handle = _BytesIOWrapper(
947 handle,
948 encoding=ioargs.encoding,
949 )
950 elif is_text and (
951 compression or memory_map or _is_binary_mode(handle, ioargs.mode)
952 ):
953 if (
954 not hasattr(handle, "readable")
955 or not hasattr(handle, "writable")
956 or not hasattr(handle, "seekable")
957 ):
958 handle = _IOWrapper(handle)
959 # error: Value of type variable "_BufferT_co" of "TextIOWrapper" cannot
960 # be "_IOWrapper | BaseBuffer" [type-var]
961 handle = TextIOWrapper(
962 handle, # type: ignore[type-var]
963 encoding=ioargs.encoding,
964 errors=errors,
965 newline="",
966 )
967 handles.append(handle)
968 # only marked as wrapped when the caller provided a handle
969 is_wrapped = not (
970 isinstance(ioargs.filepath_or_buffer, str) or ioargs.should_close
971 )
973 if "r" in ioargs.mode and not hasattr(handle, "read"):
974 raise TypeError(
975 "Expected file path name or file-like object, "
976 f"got {type(ioargs.filepath_or_buffer)} type"
977 )
979 handles.reverse() # close the most recently added buffer first
980 if ioargs.should_close:
981 assert not isinstance(ioargs.filepath_or_buffer, str)
982 handles.append(ioargs.filepath_or_buffer)
983 handles.extend(ioargs.close_handles)
985 return IOHandles(
986 # error: Argument "handle" to "IOHandles" has incompatible type
987 # "Union[TextIOWrapper, GzipFile, BaseBuffer, typing.IO[bytes],
988 # typing.IO[Any]]"; expected "pandas._typing.IO[Any]"
989 handle=handle, # type: ignore[arg-type]
990 # error: Argument "created_handles" to "IOHandles" has incompatible type
991 # "List[BaseBuffer]"; expected "List[Union[IO[bytes], IO[str]]]"
992 created_handles=handles, # type: ignore[arg-type]
993 is_wrapped=is_wrapped,
994 compression=ioargs.compression,
995 )
998class _BufferedWriter(BytesIO, ABC):
999 """
1000 Some objects do not support multiple .write() calls (TarFile and ZipFile).
1001 This wrapper writes to the underlying buffer on close.
1002 """
1004 buffer = BytesIO()
1006 @abstractmethod
1007 def write_to_buffer(self) -> None: ...
1009 def close(self) -> None:
1010 if self.closed:
1011 # already closed
1012 return
1013 if self.getbuffer().nbytes:
1014 # write to buffer
1015 self.seek(0)
1016 with self.buffer:
1017 self.write_to_buffer()
1018 else:
1019 self.buffer.close()
1020 super().close()
1023class _BytesTarFile(_BufferedWriter):
1024 def __init__(
1025 self,
1026 name: str | None = None,
1027 mode: Literal["r", "a", "w", "x"] = "r",
1028 fileobj: ReadBuffer[bytes] | WriteBuffer[bytes] | None = None,
1029 archive_name: str | None = None,
1030 **kwargs: Any,
1031 ) -> None:
1032 super().__init__()
1033 self.archive_name = archive_name
1034 self.name = name
1035 # error: No overload variant of "open" of "TarFile" matches argument
1036 # types "str | None", "str", "ReadBuffer[bytes] | WriteBuffer[bytes] | None",
1037 # "dict[str, Any]"
1038 # error: Incompatible types in assignment (expression has type "TarFile",
1039 # base class "_BufferedWriter" defined the type as "BytesIO")
1040 self.buffer: tarfile.TarFile = tarfile.TarFile.open( # type: ignore[call-overload, assignment]
1041 name=name,
1042 mode=self.extend_mode(mode),
1043 fileobj=fileobj,
1044 **kwargs,
1045 )
1047 def extend_mode(self, mode: str) -> str:
1048 mode = mode.replace("b", "")
1049 if mode != "w":
1050 return mode
1051 if self.name is not None:
1052 suffix = Path(self.name).suffix
1053 if suffix in (".gz", ".xz", ".bz2"):
1054 mode = f"{mode}:{suffix[1:]}"
1055 return mode
1057 def infer_filename(self) -> str | None:
1058 """
1059 If an explicit archive_name is not given, we still want the file inside the zip
1060 file not to be named something.tar, because that causes confusion (GH39465).
1061 """
1062 if self.name is None:
1063 return None
1065 filename = Path(self.name)
1066 if filename.suffix == ".tar":
1067 return filename.with_suffix("").name
1068 elif filename.suffix in (".tar.gz", ".tar.bz2", ".tar.xz"):
1069 return filename.with_suffix("").with_suffix("").name
1070 return filename.name
1072 def write_to_buffer(self) -> None:
1073 # TarFile needs a non-empty string
1074 archive_name = self.archive_name or self.infer_filename() or "tar"
1075 tarinfo = tarfile.TarInfo(name=archive_name)
1076 tarinfo.size = len(self.getvalue())
1077 self.buffer.addfile(tarinfo, self)
1080class _BytesZipFile(_BufferedWriter):
1081 def __init__(
1082 self,
1083 file: FilePath | ReadBuffer[bytes] | WriteBuffer[bytes],
1084 mode: str,
1085 archive_name: str | None = None,
1086 **kwargs: Any,
1087 ) -> None:
1088 super().__init__()
1089 mode = mode.replace("b", "")
1090 self.archive_name = archive_name
1092 kwargs.setdefault("compression", zipfile.ZIP_DEFLATED)
1093 # error: No overload variant of "ZipFile" matches argument types
1094 # "str | PathLike[str] | ReadBuffer[bytes] | WriteBuffer[bytes]",
1095 # "str", "dict[str, Any]"
1096 # error: Incompatible types in assignment (expression has type "ZipFile",
1097 # base class "_BufferedWriter" defined the type as "BytesIO")
1098 self.buffer: zipfile.ZipFile = zipfile.ZipFile( # type: ignore[call-overload, assignment]
1099 file, mode, **kwargs
1100 )
1102 def infer_filename(self) -> str | None:
1103 """
1104 If an explicit archive_name is not given, we still want the file inside the zip
1105 file not to be named something.zip, because that causes confusion (GH39465).
1106 """
1107 if isinstance(self.buffer.filename, (os.PathLike, str)):
1108 filename = Path(self.buffer.filename)
1109 if filename.suffix == ".zip":
1110 return filename.with_suffix("").name
1111 return filename.name
1112 return None
1114 def write_to_buffer(self) -> None:
1115 # ZipFile needs a non-empty string
1116 archive_name = self.archive_name or self.infer_filename() or "zip"
1117 self.buffer.writestr(archive_name, self.getvalue())
1120class _IOWrapper:
1121 # TextIOWrapper is overly strict: it request that the buffer has seekable, readable,
1122 # and writable. If we have a read-only buffer, we shouldn't need writable and vice
1123 # versa. Some buffers, are seek/read/writ-able but they do not have the "-able"
1124 # methods, e.g., tempfile.SpooledTemporaryFile.
1125 # If a buffer does not have the above "-able" methods, we simple assume they are
1126 # seek/read/writ-able.
1127 def __init__(self, buffer: BaseBuffer) -> None:
1128 self.buffer = buffer
1130 def __getattr__(self, name: str) -> Any:
1131 return getattr(self.buffer, name)
1133 def readable(self) -> bool:
1134 if hasattr(self.buffer, "readable"):
1135 return self.buffer.readable()
1136 return True
1138 def seekable(self) -> bool:
1139 if hasattr(self.buffer, "seekable"):
1140 return self.buffer.seekable()
1141 return True
1143 def writable(self) -> bool:
1144 if hasattr(self.buffer, "writable"):
1145 return self.buffer.writable()
1146 return True
1149class _BytesIOWrapper:
1150 # Wrapper that wraps a StringIO buffer and reads bytes from it
1151 # Created for compat with pyarrow read_csv
1152 def __init__(self, buffer: StringIO | TextIOBase, encoding: str = "utf-8") -> None:
1153 self.buffer = buffer
1154 self.encoding = encoding
1155 # Because a character can be represented by more than 1 byte,
1156 # it is possible that reading will produce more bytes than n
1157 # We store the extra bytes in this overflow variable, and append the
1158 # overflow to the front of the bytestring the next time reading is performed
1159 self.overflow = b""
1161 def __getattr__(self, attr: str) -> Any:
1162 return getattr(self.buffer, attr)
1164 def read(self, n: int | None = -1) -> bytes:
1165 assert self.buffer is not None
1166 bytestring = self.buffer.read(n).encode(self.encoding)
1167 # When n=-1/n greater than remaining bytes: Read entire file/rest of file
1168 combined_bytestring = self.overflow + bytestring
1169 if n is None or n < 0 or n >= len(combined_bytestring):
1170 self.overflow = b""
1171 return combined_bytestring
1172 else:
1173 to_return = combined_bytestring[:n]
1174 self.overflow = combined_bytestring[n:]
1175 return to_return
1178def _maybe_memory_map(
1179 handle: str | BaseBuffer, memory_map: bool
1180) -> tuple[str | BaseBuffer, bool, list[BaseBuffer]]:
1181 """Try to memory map file/buffer."""
1182 handles: list[BaseBuffer] = []
1183 memory_map &= hasattr(handle, "fileno") or isinstance(handle, str)
1184 if not memory_map:
1185 return handle, memory_map, handles
1187 # mmap used by only read_csv
1188 handle = cast(ReadCsvBuffer, handle)
1190 # need to open the file first
1191 if isinstance(handle, str):
1192 handle = open(handle, "rb")
1193 handles.append(handle)
1195 try:
1196 # open mmap and adds *-able
1197 # error: Argument 1 to "_IOWrapper" has incompatible type "mmap";
1198 # expected "BaseBuffer"
1199 wrapped = _IOWrapper(
1200 mmap.mmap(
1201 handle.fileno(),
1202 0,
1203 access=mmap.ACCESS_READ, # type: ignore[arg-type]
1204 )
1205 )
1206 finally:
1207 for handle in reversed(handles):
1208 # error: "BaseBuffer" has no attribute "close"
1209 handle.close() # type: ignore[attr-defined]
1211 return wrapped, memory_map, [wrapped]
1214def file_exists(filepath_or_buffer: FilePath | BaseBuffer) -> bool:
1215 """Test whether file exists."""
1216 exists = False
1217 filepath_or_buffer = stringify_path(filepath_or_buffer)
1218 if not isinstance(filepath_or_buffer, str):
1219 return exists
1220 try:
1221 exists = os.path.exists(filepath_or_buffer)
1222 # gh-5874: if the filepath is too long will raise here
1223 except (TypeError, ValueError):
1224 pass
1225 return exists
1228def _is_binary_mode(handle: FilePath | BaseBuffer, mode: str) -> bool:
1229 """Whether the handle is opened in binary mode"""
1230 # specified by user
1231 if "t" in mode or "b" in mode:
1232 return "b" in mode
1234 # exceptions
1235 text_classes = (
1236 # classes that expect string but have 'b' in mode
1237 codecs.StreamWriter,
1238 codecs.StreamReader,
1239 codecs.StreamReaderWriter,
1240 )
1241 if issubclass(type(handle), text_classes):
1242 return False
1244 return isinstance(handle, _get_binary_io_classes()) or "b" in getattr(
1245 handle, "mode", mode
1246 )
1249@functools.lru_cache
1250def _get_binary_io_classes() -> tuple[type, ...]:
1251 """IO classes that that expect bytes"""
1252 binary_classes: tuple[type, ...] = (BufferedIOBase, RawIOBase)
1254 # python-zstandard doesn't use any of the builtin base classes; instead we
1255 # have to use the `zstd.ZstdDecompressionReader` class for isinstance checks.
1256 # Unfortunately `zstd.ZstdDecompressionReader` isn't exposed by python-zstandard
1257 # so we have to get it from a `zstd.ZstdDecompressor` instance.
1258 # See also https://github.com/indygreg/python-zstandard/pull/165.
1259 zstd = import_optional_dependency("zstandard", errors="ignore")
1260 if zstd is not None:
1261 with zstd.ZstdDecompressor().stream_reader(b"") as reader:
1262 binary_classes += (type(reader),)
1264 return binary_classes
1267def is_potential_multi_index(
1268 columns: Sequence[Hashable] | MultiIndex,
1269 index_col: bool | Sequence[int] | None = None,
1270) -> bool:
1271 """
1272 Check whether or not the `columns` parameter
1273 could be converted into a MultiIndex.
1275 Parameters
1276 ----------
1277 columns : array-like
1278 Object which may or may not be convertible into a MultiIndex
1279 index_col : None, bool or list, optional
1280 Column or columns to use as the (possibly hierarchical) index
1282 Returns
1283 -------
1284 bool : Whether or not columns could become a MultiIndex
1285 """
1286 if index_col is None or isinstance(index_col, bool):
1287 index_columns = set()
1288 else:
1289 index_columns = set(index_col)
1291 return bool(
1292 len(columns)
1293 and not isinstance(columns, ABCMultiIndex)
1294 and all(isinstance(c, tuple) for c in columns if c not in index_columns)
1295 )
1298def dedup_names(
1299 names: Sequence[Hashable], is_potential_multiindex: bool
1300) -> Sequence[Hashable]:
1301 """
1302 Rename column names if duplicates exist.
1304 Currently the renaming is done by appending a period and an autonumeric,
1305 but a custom pattern may be supported in the future.
1307 Examples
1308 --------
1309 >>> dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False)
1310 ['x', 'y', 'x.1', 'x.2']
1311 """
1312 names = list(names) # so we can index
1313 counts: DefaultDict[Hashable, int] = defaultdict(int)
1315 for i, col in enumerate(names):
1316 cur_count = counts[col]
1318 while cur_count > 0:
1319 counts[col] = cur_count + 1
1321 if is_potential_multiindex:
1322 # for mypy
1323 assert isinstance(col, tuple)
1324 col = (*col[:-1], f"{col[-1]}.{cur_count}")
1325 else:
1326 col = f"{col}.{cur_count}"
1327 cur_count = counts[col]
1329 names[i] = col
1330 counts[col] = cur_count + 1
1332 return names