"""Common IO api utilities"""
from __future__ import annotations
from abc import (
ABC,
abstractmethod,
)
import codecs
from collections import defaultdict
import dataclasses
import functools
import gzip
from io import (
BufferedIOBase,
BytesIO,
RawIOBase,
StringIO,
TextIOBase,
TextIOWrapper,
)
import mmap
import os
from pathlib import Path
import re
import tarfile
from typing import (
IO,
Any,
AnyStr,
DefaultDict,
Generic,
Hashable,
Literal,
Mapping,
Sequence,
TypeVar,
cast,
overload,
)
from urllib.parse import (
urljoin,
urlparse as parse_url,
uses_netloc,
uses_params,
uses_relative,
)
import warnings
import zipfile
from pandas._typing import (
BaseBuffer,
CompressionDict,
CompressionOptions,
FilePath,
ReadBuffer,
ReadCsvBuffer,
StorageOptions,
WriteBuffer,
)
from pandas.compat import get_lzma_file
from pandas.compat._optional import import_optional_dependency
from pandas.compat.compressors import BZ2File as _BZ2File
from pandas.util._decorators import doc
from pandas.util._exceptions import find_stack_level
from pandas.core.dtypes.common import (
is_bool,
is_file_like,
is_integer,
is_list_like,
)
from pandas.core.indexes.api import MultiIndex
from pandas.core.shared_docs import _shared_docs
_VALID_URLS = set(uses_relative + uses_netloc + uses_params)
_VALID_URLS.discard("")
_RFC_3986_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9+\-+.]*://")
BaseBufferT = TypeVar("BaseBufferT", bound=BaseBuffer)
@dataclasses.dataclass
class IOArgs:
"""
Return value of io/common.py:_get_filepath_or_buffer.
"""
filepath_or_buffer: str | BaseBuffer
encoding: str
mode: str
compression: CompressionDict
should_close: bool = False
@dataclasses.dataclass
class IOHandles(Generic[AnyStr]):
"""
Return value of io/common.py:get_handle
Can be used as a context manager.
This is used to easily close created buffers and to handle corner cases when
TextIOWrapper is inserted.
handle: The file handle to be used.
created_handles: All file handles that are created by get_handle
is_wrapped: Whether a TextIOWrapper needs to be detached.
"""
# handle might not implement the IO-interface
handle: IO[AnyStr]
compression: CompressionDict
created_handles: list[IO[bytes] | IO[str]] = dataclasses.field(default_factory=list)
is_wrapped: bool = False
def close(self) -> None:
"""
Close all created buffers.
Note: If a TextIOWrapper was inserted, it is flushed and detached to
avoid closing the potentially user-created buffer.
"""
if self.is_wrapped:
assert isinstance(self.handle, TextIOWrapper)
self.handle.flush()
self.handle.detach()
self.created_handles.remove(self.handle)
for handle in self.created_handles:
handle.close()
self.created_handles = []
self.is_wrapped = False
def __enter__(self) -> IOHandles[AnyStr]:
return self
def __exit__(self, *args: Any) -> None:
self.close()
def is_url(url: object) -> bool:
"""
Check to see if a URL has a valid protocol.
Parameters
----------
url : str or unicode
Returns
-------
isurl : bool
If `url` has a valid protocol return True otherwise False.
"""
if not isinstance(url, str):
return False
return parse_url(url).scheme in _VALID_URLS
@overload
def _expand_user(filepath_or_buffer: str) -> str:
...
@overload
def _expand_user(filepath_or_buffer: BaseBufferT) -> BaseBufferT:
...
def _expand_user(filepath_or_buffer: str | BaseBufferT) -> str | BaseBufferT:
"""
Return the argument with an initial component of ~ or ~user
replaced by that user's home directory.
Parameters
----------
filepath_or_buffer : object to be converted if possible
Returns
-------
expanded_filepath_or_buffer : an expanded filepath or the
input if not expandable
"""
if isinstance(filepath_or_buffer, str):
return os.path.expanduser(filepath_or_buffer)
return filepath_or_buffer
def validate_header_arg(header: object) -> None:
if header is None:
return
if is_integer(header):
header = cast(int, header)
if header < 0:
# GH 27779
raise ValueError(
"Passing negative integer to header is invalid. "
"For no header, use header=None instead"
)
return
if is_list_like(header, allow_sets=False):
header = cast(Sequence, header)
if not all(map(is_integer, header)):
raise ValueError("header must be integer or list of integers")
if any(i < 0 for i in header):
raise ValueError("cannot specify multi-index header with negative integers")
return
if is_bool(header):
raise TypeError(
"Passing a bool to header is invalid. Use header=None for no header or "
"header=int or list-like of ints to specify "
"the row(s) making up the column names"
)
# GH 16338
raise ValueError("header must be integer or list of integers")
@overload
def stringify_path(filepath_or_buffer: FilePath, convert_file_like: bool = ...) -> str:
...
@overload
def stringify_path(
filepath_or_buffer: BaseBufferT, convert_file_like: bool = ...
) -> BaseBufferT:
...
def stringify_path(
filepath_or_buffer: FilePath | BaseBufferT,
convert_file_like: bool = False,
) -> str | BaseBufferT:
"""
Attempt to convert a path-like object to a string.
Parameters
----------
filepath_or_buffer : object to be converted
Returns
-------
str_filepath_or_buffer : maybe a string version of the object
Notes
-----
Objects supporting the fspath protocol (python 3.6+) are coerced
according to its __fspath__ method.
Any other object is passed through unchanged, which includes bytes,
strings, buffers, or anything else that's not even path-like.
"""
if not convert_file_like and is_file_like(filepath_or_buffer):
# GH 38125: some fsspec objects implement os.PathLike but have already opened a
# file. This prevents opening the file a second time. infer_compression calls
# this function with convert_file_like=True to infer the compression.
return cast(BaseBufferT, filepath_or_buffer)
if isinstance(filepath_or_buffer, os.PathLike):
filepath_or_buffer = filepath_or_buffer.__fspath__()
return _expand_user(filepath_or_buffer)
def urlopen(*args, **kwargs):
"""
Lazy-import wrapper for stdlib urlopen, as that imports a big chunk of
the stdlib.
"""
import urllib.request
return urllib.request.urlopen(*args, **kwargs)
def is_fsspec_url(url: FilePath | BaseBuffer) -> bool:
"""
Returns true if the given URL looks like
something fsspec can handle
"""
return (
isinstance(url, str)
and bool(_RFC_3986_PATTERN.match(url))
and not url.startswith(("http://", "https://"))
)
@doc(
storage_options=_shared_docs["storage_options"],
compression_options=_shared_docs["compression_options"] % "filepath_or_buffer",
)
def _get_filepath_or_buffer(
filepath_or_buffer: FilePath | BaseBuffer,
encoding: str = "utf-8",
compression: CompressionOptions = None,
mode: str = "r",
storage_options: StorageOptions = None,
) -> IOArgs:
"""
If the filepath_or_buffer is a url, translate and return the buffer.
Otherwise passthrough.
Parameters
----------
filepath_or_buffer : a url, filepath (str, py.path.local or pathlib.Path),
or buffer
{compression_options}
.. versionchanged:: 1.4.0 Zstandard support.
encoding : the encoding to use to decode bytes, default is 'utf-8'
mode : str, optional
{storage_options}
.. versionadded:: 1.2.0
..versionchange:: 1.2.0
Returns the dataclass IOArgs.
"""
filepath_or_buffer = stringify_path(filepath_or_buffer)
# handle compression dict
compression_method, compression = get_compression_method(compression)
compression_method = infer_compression(filepath_or_buffer, compression_method)
# GH21227 internal compression is not used for non-binary handles.
if compression_method and hasattr(filepath_or_buffer, "write") and "b" not in mode:
warnings.warn(
"compression has no effect when passing a non-binary object as input.",
RuntimeWarning,
stacklevel=find_stack_level(),
)
compression_method = None
compression = dict(compression, method=compression_method)
# bz2 and xz do not write the byte order mark for utf-16 and utf-32
# print a warning when writing such files
if (
"w" in mode
and compression_method in ["bz2", "xz"]
and encoding in ["utf-16", "utf-32"]
):
warnings.warn(
f"{compression} will not write the byte order mark for {encoding}",
UnicodeWarning,
stacklevel=find_stack_level(),
)
# Use binary mode when converting path-like objects to file-like objects (fsspec)
# except when text mode is explicitly requested. The original mode is returned if
# fsspec is not used.
fsspec_mode = mode
if "t" not in fsspec_mode and "b" not in fsspec_mode:
fsspec_mode += "b"
if isinstance(filepath_or_buffer, str) and is_url(filepath_or_buffer):
# TODO: fsspec can also handle HTTP via requests, but leaving this
# unchanged. using fsspec appears to break the ability to infer if the
# server responded with gzipped data
storage_options = storage_options or {}
# waiting until now for importing to match intended lazy logic of
# urlopen function defined elsewhere in this module
import urllib.request
# assuming storage_options is to be interpreted as headers
req_info = urllib.request.Request(filepath_or_buffer, headers=storage_options)
with urlopen(req_info) as req:
content_encoding = req.headers.get("Content-Encoding", None)
if content_encoding == "gzip":
# Override compression based on Content-Encoding header
compression = {"method": "gzip"}
reader = BytesIO(req.read())
return IOArgs(
filepath_or_buffer=reader,
encoding=encoding,
compression=compression,
should_close=True,
mode=fsspec_mode,
)
if is_fsspec_url(filepath_or_buffer):
assert isinstance(
filepath_or_buffer, str
) # just to appease mypy for this branch
# two special-case s3-like protocols; these have special meaning in Hadoop,
# but are equivalent to just "s3" from fsspec's point of view
# cc #11071
if filepath_or_buffer.startswith("s3a://"):
filepath_or_buffer = filepath_or_buffer.replace("s3a://", "s3://")
if filepath_or_buffer.startswith("s3n://"):
filepath_or_buffer = filepath_or_buffer.replace("s3n://", "s3://")
fsspec = import_optional_dependency("fsspec")
# If botocore is installed we fallback to reading with anon=True
# to allow reads from public buckets
err_types_to_retry_with_anon: list[Any] = []
try:
import_optional_dependency("botocore")
from botocore.exceptions import (
ClientError,
NoCredentialsError,
)
err_types_to_retry_with_anon = [
ClientError,
NoCredentialsError,
PermissionError,
]
except ImportError:
pass
try:
file_obj = fsspec.open(
filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
).open()
# GH 34626 Reads from Public Buckets without Credentials needs anon=True
except tuple(err_types_to_retry_with_anon):
if storage_options is None:
storage_options = {"anon": True}
else:
# don't mutate user input.
storage_options = dict(storage_options)
storage_options["anon"] = True
file_obj = fsspec.open(
filepath_or_buffer, mode=fsspec_mode, **(storage_options or {})
).open()
return IOArgs(
filepath_or_buffer=file_obj,
encoding=encoding,
compression=compression,
should_close=True,
mode=fsspec_mode,
)
elif storage_options:
raise ValueError(
"storage_options passed with file object or non-fsspec file path"
)
if isinstance(filepath_or_buffer, (str, bytes, mmap.mmap)):
return IOArgs(
filepath_or_buffer=_expand_user(filepath_or_buffer),
encoding=encoding,
compression=compression,
should_close=False,
mode=mode,
)
# is_file_like requires (read | write) & __iter__ but __iter__ is only
# needed for read_csv(engine=python)
if not (
hasattr(filepath_or_buffer, "read") or hasattr(filepath_or_buffer, "write")
):
msg = f"Invalid file path or buffer object type: {type(filepath_or_buffer)}"
raise ValueError(msg)
return IOArgs(
filepath_or_buffer=filepath_or_buffer,
encoding=encoding,
compression=compression,
should_close=False,
mode=mode,
)
def file_path_to_url(path: str) -> str:
"""
converts an absolute native path to a FILE URL.
Parameters
----------
path : a path in native format
Returns
-------
a valid FILE URL
"""
# lazify expensive import (~30ms)
from urllib.request import pathname2url
return urljoin("file:", pathname2url(path))
extension_to_compression = {
".tar": "tar",
".tar.gz": "tar",
".tar.bz2": "tar",
".tar.xz": "tar",
".gz": "gzip",
".bz2": "bz2",
".zip": "zip",
".xz": "xz",
".zst": "zstd",
}
_supported_compressions = set(extension_to_compression.values())
def get_compression_method(
compression: CompressionOptions,
) -> tuple[str | None, CompressionDict]:
"""
Simplifies a compression argument to a compression method string and
a mapping containing additional arguments.
Parameters
----------
compression : str or mapping
If string, specifies the compression method. If mapping, value at key
'method' specifies compression method.
Returns
-------
tuple of ({compression method}, Optional[str]
{compression arguments}, Dict[str, Any])
Raises
------
ValueError on mapping missing 'method' key
"""
compression_method: str | None
if isinstance(compression, Mapping):
compression_args = dict(compression)
try:
compression_method = compression_args.pop("method")
except KeyError as err:
raise ValueError("If mapping, compression must have key 'method'") from err
else:
compression_args = {}
compression_method = compression
return compression_method, compression_args
@doc(compression_options=_shared_docs["compression_options"] % "filepath_or_buffer")
def infer_compression(
filepath_or_buffer: FilePath | BaseBuffer, compression: str | None
) -> str | None:
"""
Get the compression method for filepath_or_buffer. If compression='infer',
the inferred compression method is returned. Otherwise, the input
compression method is returned unchanged, unless it's invalid, in which
case an error is raised.
Parameters
----------
filepath_or_buffer : str or file handle
File path or object.
{compression_options}
.. versionchanged:: 1.4.0 Zstandard support.
Returns
-------
string or None
Raises
------
ValueError on invalid compression specified.
"""
if compression is None:
return None
# Infer compression
if compression == "infer":
# Convert all path types (e.g. pathlib.Path) to strings
filepath_or_buffer = stringify_path(filepath_or_buffer, convert_file_like=True)
if not isinstance(filepath_or_buffer, str):
# Cannot infer compression of a buffer, assume no compression
return None
# Infer compression from the filename/URL extension
for extension, compression in extension_to_compression.items():
if filepath_or_buffer.lower().endswith(extension):
return compression
return None
# Compression has been specified. Check that it's valid
if compression in _supported_compressions:
return compression
valid = ["infer", None] + sorted(_supported_compressions)
msg = (
f"Unrecognized compression type: {compression}\n"
f"Valid compression types are {valid}"
)
raise ValueError(msg)
def check_parent_directory(path: Path | str) -> None:
"""
Check if parent directory of a file exists, raise OSError if it does not
Parameters
----------
path: Path or str
Path to check parent directory of
"""
parent = Path(path).parent
if not parent.is_dir():
raise OSError(rf"Cannot save file into a non-existent directory: '{parent}'")
@overload
def get_handle(
path_or_buf: FilePath | BaseBuffer,
mode: str,
*,
encoding: str | None = ...,
compression: CompressionOptions = ...,
memory_map: bool = ...,
is_text: Literal[False],
errors: str | None = ...,
storage_options: StorageOptions = ...,
) -> IOHandles[bytes]:
...
@overload
def get_handle(
path_or_buf: FilePath | BaseBuffer,
mode: str,
*,
encoding: str | None = ...,
compression: CompressionOptions = ...,
memory_map: bool = ...,
is_text: Literal[True] = ...,
errors: str | None = ...,
storage_options: StorageOptions = ...,
) -> IOHandles[str]:
...
@overload
def get_handle(
path_or_buf: FilePath | BaseBuffer,
mode: str,
*,
encoding: str | None = ...,
compression: CompressionOptions = ...,
memory_map: bool = ...,
is_text: bool = ...,
errors: str | None = ...,
storage_options: StorageOptions = ...,
) -> IOHandles[str] | IOHandles[bytes]:
...
@doc(compression_options=_shared_docs["compression_options"] % "path_or_buf")
def get_handle(
path_or_buf: FilePath | BaseBuffer,
mode: str,
*,
encoding: str | None = None,
compression: CompressionOptions = None,
memory_map: bool = False,
is_text: bool = True,
errors: str | None = None,
storage_options: StorageOptions = None,
) -> IOHandles[str] | IOHandles[bytes]:
"""
Get file handle for given path/buffer and mode.
Parameters
----------
path_or_buf : str or file handle
File path or object.
mode : str
Mode to open path_or_buf with.
encoding : str or None
Encoding to use.
{compression_options}
.. versionchanged:: 1.0.0
May now be a dict with key 'method' as compression mode
and other keys as compression options if compression
mode is 'zip'.
.. versionchanged:: 1.1.0
Passing compression options as keys in dict is now
supported for compression modes 'gzip', 'bz2', 'zstd' and 'zip'.
.. versionchanged:: 1.4.0 Zstandard support.
memory_map : bool, default False
See parsers._parser_params for more information. Only used by read_csv.
is_text : bool, default True
Whether the type of the content passed to the file/buffer is string or
bytes. This is not the same as `"b" not in mode`. If a string content is
passed to a binary file/buffer, a wrapper is inserted.
errors : str, default 'strict'
Specifies how encoding and decoding errors are to be handled.
See the errors argument for :func:`open` for a full list
of options.
storage_options: StorageOptions = None
Passed to _get_filepath_or_buffer
.. versionchanged:: 1.2.0
Returns the dataclass IOHandles
"""
# Windows does not default to utf-8. Set to utf-8 for a consistent behavior
encoding = encoding or "utf-8"
errors = errors or "strict"
# read_csv does not know whether the buffer is opened in binary/text mode
if _is_binary_mode(path_or_buf, mode) and "b" not in mode:
mode += "b"
# validate encoding and errors
codecs.lookup(encoding)
if isinstance(errors, str):
codecs.lookup_error(errors)
# open URLs
ioargs = _get_filepath_or_buffer(
path_or_buf,
encoding=encoding,
compression=compression,
mode=mode,
storage_options=storage_options,
)
handle = ioargs.filepath_or_buffer
handles: list[BaseBuffer]
# memory mapping needs to be the first step
# only used for read_csv
handle, memory_map, handles = _maybe_memory_map(handle, memory_map)
is_path = isinstance(handle, str)
compression_args = dict(ioargs.compression)
compression = compression_args.pop("method")
# Only for write methods
if "r" not in mode and is_path:
check_parent_directory(str(handle))
if compression:
if compression != "zstd":
# compression libraries do not like an explicit text-mode
ioargs.mode = ioargs.mode.replace("t", "")
elif compression == "zstd" and "b" not in ioargs.mode:
# python-zstandard defaults to text mode, but we always expect
# compression libraries to use binary mode.
ioargs.mode += "b"
# GZ Compression
if compression == "gzip":
if isinstance(handle, str):
# error: Incompatible types in assignment (expression has type
# "GzipFile", variable has type "Union[str, BaseBuffer]")
handle = gzip.GzipFile( # type: ignore[assignment]
filename=handle,
mode=ioargs.mode,
**compression_args,
)
else:
handle = gzip.GzipFile(
# No overload variant of "GzipFile" matches argument types
# "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
fileobj=handle, # type: ignore[call-overload]
mode=ioargs.mode,
**compression_args,
)
# BZ Compression
elif compression == "bz2":
# Overload of "BZ2File" to handle pickle protocol 5
# "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
handle = _BZ2File( # type: ignore[call-overload]
handle,
mode=ioargs.mode,
**compression_args,
)
# ZIP Compression
elif compression == "zip":
# error: Argument 1 to "_BytesZipFile" has incompatible type
# "Union[str, BaseBuffer]"; expected "Union[Union[str, PathLike[str]],
# ReadBuffer[bytes], WriteBuffer[bytes]]"
handle = _BytesZipFile(
handle, ioargs.mode, **compression_args # type: ignore[arg-type]
)
if handle.buffer.mode == "r":
handles.append(handle)
zip_names = handle.buffer.namelist()
if len(zip_names) == 1:
handle = handle.buffer.open(zip_names.pop())
elif not zip_names:
raise ValueError(f"Zero files found in ZIP file {path_or_buf}")
else:
raise ValueError(
"Multiple files found in ZIP file. "
f"Only one file per ZIP: {zip_names}"
)
# TAR Encoding
elif compression == "tar":
compression_args.setdefault("mode", ioargs.mode)
if isinstance(handle, str):
handle = _BytesTarFile(name=handle, **compression_args)
else:
# error: Argument "fileobj" to "_BytesTarFile" has incompatible
# type "BaseBuffer"; expected "Union[ReadBuffer[bytes],
# WriteBuffer[bytes], None]"
handle = _BytesTarFile(
fileobj=handle, **compression_args # type: ignore[arg-type]
)
assert isinstance(handle, _BytesTarFile)
if "r" in handle.buffer.mode:
handles.append(handle)
files = handle.buffer.getnames()
if len(files) == 1:
file = handle.buffer.extractfile(files[0])
assert file is not None
handle = file
elif not files:
raise ValueError(f"Zero files found in TAR archive {path_or_buf}")
else:
raise ValueError(
"Multiple files found in TAR archive. "
f"Only one file per TAR archive: {files}"
)
# XZ Compression
elif compression == "xz":
# error: Argument 1 to "LZMAFile" has incompatible type "Union[str,
# BaseBuffer]"; expected "Optional[Union[Union[str, bytes, PathLike[str],
# PathLike[bytes]], IO[bytes]]]"
handle = get_lzma_file()(handle, ioargs.mode) # type: ignore[arg-type]
# Zstd Compression
elif compression == "zstd":
zstd = import_optional_dependency("zstandard")
if "r" in ioargs.mode:
open_args = {"dctx": zstd.ZstdDecompressor(**compression_args)}
else:
open_args = {"cctx": zstd.ZstdCompressor(**compression_args)}
handle = zstd.open(
handle,
mode=ioargs.mode,
**open_args,
)
# Unrecognized Compression
else:
msg = f"Unrecognized compression type: {compression}"
raise ValueError(msg)
assert not isinstance(handle, str)
handles.append(handle)
elif isinstance(handle, str):
# Check whether the filename is to be opened in binary mode.
# Binary mode does not support 'encoding' and 'newline'.
if ioargs.encoding and "b" not in ioargs.mode:
# Encoding
handle = open(
handle,
ioargs.mode,
encoding=ioargs.encoding,
errors=errors,
newline="",
)
else:
# Binary mode
handle = open(handle, ioargs.mode)
handles.append(handle)
# Convert BytesIO or file objects passed with an encoding
is_wrapped = False
if not is_text and ioargs.mode == "rb" and isinstance(handle, TextIOBase):
# not added to handles as it does not open/buffer resources
handle = _BytesIOWrapper(
handle,
encoding=ioargs.encoding,
)
elif is_text and (
compression or memory_map or _is_binary_mode(handle, ioargs.mode)
):
if (
not hasattr(handle, "readable")
or not hasattr(handle, "writable")
or not hasattr(handle, "seekable")
):
handle = _IOWrapper(handle)
# error: Argument 1 to "TextIOWrapper" has incompatible type
# "_IOWrapper"; expected "IO[bytes]"
handle = TextIOWrapper(
handle, # type: ignore[arg-type]
encoding=ioargs.encoding,
errors=errors,
newline="",
)
handles.append(handle)
# only marked as wrapped when the caller provided a handle
is_wrapped = not (
isinstance(ioargs.filepath_or_buffer, str) or ioargs.should_close
)
if "r" in ioargs.mode and not hasattr(handle, "read"):
raise TypeError(
"Expected file path name or file-like object, "
f"got {type(ioargs.filepath_or_buffer)} type"
)
handles.reverse() # close the most recently added buffer first
if ioargs.should_close:
assert not isinstance(ioargs.filepath_or_buffer, str)
handles.append(ioargs.filepath_or_buffer)
return IOHandles(
# error: Argument "handle" to "IOHandles" has incompatible type
# "Union[TextIOWrapper, GzipFile, BaseBuffer, typing.IO[bytes],
# typing.IO[Any]]"; expected "pandas._typing.IO[Any]"
handle=handle, # type: ignore[arg-type]
# error: Argument "created_handles" to "IOHandles" has incompatible type
# "List[BaseBuffer]"; expected "List[Union[IO[bytes], IO[str]]]"
created_handles=handles, # type: ignore[arg-type]
is_wrapped=is_wrapped,
compression=ioargs.compression,
)
# error: Definition of "__enter__" in base class "IOBase" is incompatible
# with definition in base class "BinaryIO"
class _BufferedWriter(BytesIO, ABC): # type: ignore[misc]
"""
Some objects do not support multiple .write() calls (TarFile and ZipFile).
This wrapper writes to the underlying buffer on close.
"""
@abstractmethod
def write_to_buffer(self) -> None:
...
def close(self) -> None:
if self.closed:
# already closed
return
if self.getvalue():
# write to buffer
self.seek(0)
# error: "_BufferedWriter" has no attribute "buffer"
with self.buffer: # type: ignore[attr-defined]
self.write_to_buffer()
else:
# error: "_BufferedWriter" has no attribute "buffer"
self.buffer.close() # type: ignore[attr-defined]
super().close()
class _BytesTarFile(_BufferedWriter):
def __init__(
self,
name: str | None = None,
mode: Literal["r", "a", "w", "x"] = "r",
fileobj: ReadBuffer[bytes] | WriteBuffer[bytes] | None = None,
archive_name: str | None = None,
**kwargs,
) -> None:
super().__init__()
self.archive_name = archive_name
self.name = name
# error: Argument "fileobj" to "open" of "TarFile" has incompatible
# type "Union[ReadBuffer[bytes], WriteBuffer[bytes], None]"; expected
# "Optional[IO[bytes]]"
self.buffer = tarfile.TarFile.open(
name=name,
mode=self.extend_mode(mode),
fileobj=fileobj, # type: ignore[arg-type]
**kwargs,
)
def extend_mode(self, mode: str) -> str:
mode = mode.replace("b", "")
if mode != "w":
return mode
if self.name is not None:
suffix = Path(self.name).suffix
if suffix in (".gz", ".xz", ".bz2"):
mode = f"{mode}:{suffix[1:]}"
return mode
def infer_filename(self) -> str | None:
"""
If an explicit archive_name is not given, we still want the file inside the zip
file not to be named something.tar, because that causes confusion (GH39465).
"""
if self.name is None:
return None
filename = Path(self.name)
if filename.suffix == ".tar":
return filename.with_suffix("").name
elif filename.suffix in (".tar.gz", ".tar.bz2", ".tar.xz"):
return filename.with_suffix("").with_suffix("").name
return filename.name
def write_to_buffer(self) -> None:
# TarFile needs a non-empty string
archive_name = self.archive_name or self.infer_filename() or "tar"
tarinfo = tarfile.TarInfo(name=archive_name)
tarinfo.size = len(self.getvalue())
self.buffer.addfile(tarinfo, self)
class _BytesZipFile(_BufferedWriter):
def __init__(
self,
file: FilePath | ReadBuffer[bytes] | WriteBuffer[bytes],
mode: str,
archive_name: str | None = None,
**kwargs,
) -> None:
super().__init__()
mode = mode.replace("b", "")
self.archive_name = archive_name
kwargs.setdefault("compression", zipfile.ZIP_DEFLATED)
# error: Argument 1 to "ZipFile" has incompatible type "Union[
# Union[str, PathLike[str]], ReadBuffer[bytes], WriteBuffer[bytes]]";
# expected "Union[Union[str, PathLike[str]], IO[bytes]]"
self.buffer = zipfile.ZipFile(file, mode, **kwargs) # type: ignore[arg-type]
def infer_filename(self) -> str | None:
"""
If an explicit archive_name is not given, we still want the file inside the zip
file not to be named something.zip, because that causes confusion (GH39465).
"""
if isinstance(self.buffer.filename, (os.PathLike, str)):
filename = Path(self.buffer.filename)
if filename.suffix == ".zip":
return filename.with_suffix("").name
return filename.name
return None
def write_to_buffer(self) -> None:
# ZipFile needs a non-empty string
archive_name = self.archive_name or self.infer_filename() or "zip"
self.buffer.writestr(archive_name, self.getvalue())
class _IOWrapper:
# TextIOWrapper is overly strict: it request that the buffer has seekable, readable,
# and writable. If we have a read-only buffer, we shouldn't need writable and vice
# versa. Some buffers, are seek/read/writ-able but they do not have the "-able"
# methods, e.g., tempfile.SpooledTemporaryFile.
# If a buffer does not have the above "-able" methods, we simple assume they are
# seek/read/writ-able.
def __init__(self, buffer: BaseBuffer) -> None:
self.buffer = buffer
def __getattr__(self, name: str):
return getattr(self.buffer, name)
def readable(self) -> bool:
if hasattr(self.buffer, "readable"):
return self.buffer.readable()
return True
def seekable(self) -> bool:
if hasattr(self.buffer, "seekable"):
return self.buffer.seekable()
return True
def writable(self) -> bool:
if hasattr(self.buffer, "writable"):
return self.buffer.writable()
return True
class _BytesIOWrapper:
# Wrapper that wraps a StringIO buffer and reads bytes from it
# Created for compat with pyarrow read_csv
def __init__(self, buffer: StringIO | TextIOBase, encoding: str = "utf-8") -> None:
self.buffer = buffer
self.encoding = encoding
# Because a character can be represented by more than 1 byte,
# it is possible that reading will produce more bytes than n
# We store the extra bytes in this overflow variable, and append the
# overflow to the front of the bytestring the next time reading is performed
self.overflow = b""
def __getattr__(self, attr: str):
return getattr(self.buffer, attr)
def read(self, n: int | None = -1) -> bytes:
assert self.buffer is not None
bytestring = self.buffer.read(n).encode(self.encoding)
# When n=-1/n greater than remaining bytes: Read entire file/rest of file
combined_bytestring = self.overflow + bytestring
if n is None or n < 0 or n >= len(combined_bytestring):
self.overflow = b""
return combined_bytestring
else:
to_return = combined_bytestring[:n]
self.overflow = combined_bytestring[n:]
return to_return
def _maybe_memory_map(
handle: str | BaseBuffer, memory_map: bool
) -> tuple[str | BaseBuffer, bool, list[BaseBuffer]]:
"""Try to memory map file/buffer."""
handles: list[BaseBuffer] = []
memory_map &= hasattr(handle, "fileno") or isinstance(handle, str)
if not memory_map:
return handle, memory_map, handles
# mmap used by only read_csv
handle = cast(ReadCsvBuffer, handle)
# need to open the file first
if isinstance(handle, str):
handle = open(handle, "rb")
handles.append(handle)
try:
# open mmap and adds *-able
# error: Argument 1 to "_IOWrapper" has incompatible type "mmap";
# expected "BaseBuffer"
wrapped = _IOWrapper(
mmap.mmap(
handle.fileno(), 0, access=mmap.ACCESS_READ # type: ignore[arg-type]
)
)
finally:
for handle in reversed(handles):
# error: "BaseBuffer" has no attribute "close"
handle.close() # type: ignore[attr-defined]
return wrapped, memory_map, [wrapped]
def file_exists(filepath_or_buffer: FilePath | BaseBuffer) -> bool:
"""Test whether file exists."""
exists = False
filepath_or_buffer = stringify_path(filepath_or_buffer)
if not isinstance(filepath_or_buffer, str):
return exists
try:
exists = os.path.exists(filepath_or_buffer)
# gh-5874: if the filepath is too long will raise here
except (TypeError, ValueError):
pass
return exists
def _is_binary_mode(handle: FilePath | BaseBuffer, mode: str) -> bool:
"""Whether the handle is opened in binary mode"""
# specified by user
if "t" in mode or "b" in mode:
return "b" in mode
# exceptions
text_classes = (
# classes that expect string but have 'b' in mode
codecs.StreamWriter,
codecs.StreamReader,
codecs.StreamReaderWriter,
)
if issubclass(type(handle), text_classes):
return False
return isinstance(handle, _get_binary_io_classes()) or "b" in getattr(
handle, "mode", mode
)
@functools.lru_cache
def _get_binary_io_classes() -> tuple[type, ...]:
"""IO classes that that expect bytes"""
binary_classes: tuple[type, ...] = (BufferedIOBase, RawIOBase)
# python-zstandard doesn't use any of the builtin base classes; instead we
# have to use the `zstd.ZstdDecompressionReader` class for isinstance checks.
# Unfortunately `zstd.ZstdDecompressionReader` isn't exposed by python-zstandard
# so we have to get it from a `zstd.ZstdDecompressor` instance.
# See also https://github.com/indygreg/python-zstandard/pull/165.
zstd = import_optional_dependency("zstandard", errors="ignore")
if zstd is not None:
with zstd.ZstdDecompressor().stream_reader(b"") as reader:
binary_classes += (type(reader),)
return binary_classes
def is_potential_multi_index(
columns: Sequence[Hashable] | MultiIndex,
index_col: bool | Sequence[int] | None = None,
) -> bool:
"""
Check whether or not the `columns` parameter
could be converted into a MultiIndex.
Parameters
----------
columns : array-like
Object which may or may not be convertible into a MultiIndex
index_col : None, bool or list, optional
Column or columns to use as the (possibly hierarchical) index
Returns
-------
bool : Whether or not columns could become a MultiIndex
"""
if index_col is None or isinstance(index_col, bool):
index_col = []
return bool(
len(columns)
and not isinstance(columns, MultiIndex)
and all(isinstance(c, tuple) for c in columns if c not in list(index_col))
)
def dedup_names(
names: Sequence[Hashable], is_potential_multiindex: bool
) -> Sequence[Hashable]:
"""
Rename column names if duplicates exist.
Currently the renaming is done by appending a period and an autonumeric,
but a custom pattern may be supported in the future.
Examples
--------
>>> dedup_names(["x", "y", "x", "x"], is_potential_multiindex=False)
['x', 'y', 'x.1', 'x.2']
"""
names = list(names) # so we can index
counts: DefaultDict[Hashable, int] = defaultdict(int)
for i, col in enumerate(names):
cur_count = counts[col]
while cur_count > 0:
counts[col] = cur_count + 1
if is_potential_multiindex:
# for mypy
assert isinstance(col, tuple)
col = col[:-1] + (f"{col[-1]}.{cur_count}",)
else:
col = f"{col}.{cur_count}"
cur_count = counts[col]
names[i] = col
counts[col] = cur_count + 1
return names