r""" I/O utils (:mod:`skbio.io.util`) ================================ .. currentmodule:: skbio.io.util This module provides utility functions to deal with files and I/O in general. Functions --------- .. autosummary:: :toctree: generated/ open open_file open_files """ # ---------------------------------------------------------------------------- # Copyright (c) 2013--, scikit-bio development team. # # Distributed under the terms of the Modified BSD License. # # The full license is in the file COPYING.txt, distributed with this software. # ---------------------------------------------------------------------------- import io from contextlib import contextmanager, ExitStack from skbio.io import IOSourceError from skbio.io._iosources import get_io_sources, get_compression_handler from skbio.io._fileobject import ( is_binary_file, SaneTextIOWrapper, CompressedBufferedReader, CompressedBufferedWriter) from skbio.util._decorator import stable _d = dict(mode='r', encoding=None, errors=None, newline=None, compression='auto', compresslevel=9) def _resolve(file, mode=_d['mode'], encoding=_d['encoding'], errors=_d['errors'], newline=_d['newline'], compression=_d['compression'], compresslevel=_d['compresslevel']): arguments = locals().copy() if mode not in {'r', 'w'}: raise ValueError("Unsupported mode: %r, use 'r' or 'w'" % mode) newfile = None source = None for source_handler in get_io_sources(): source = source_handler(file, arguments) if mode == 'r' and source.can_read(): newfile = source.get_reader() break elif mode == 'w' and source.can_write(): newfile = source.get_writer() break if newfile is None: raise IOSourceError( "Could not open source: %r (mode: %r)" % (file, mode)) return newfile, source, is_binary_file(newfile) @stable(as_of="0.4.0") def open(file, mode=_d['mode'], encoding=_d['encoding'], errors=_d['errors'], newline=_d['newline'], compression=_d['compression'], compresslevel=_d['compresslevel']): r"""Convert input into a filehandle. Supported inputs: +--------------------------------------+--------+---------+-----------+ | type | can \ | can \ | source \ | | | read | write | type | +======================================+========+=========+===========+ | file path | True | True | Binary | +--------------------------------------+--------+---------+-----------+ | URL | True | False | Binary | +--------------------------------------+--------+---------+-----------+ | ``["lines list\n"]`` | True | True | Text | +--------------------------------------+--------+---------+-----------+ | :class:`io.StringIO` | True | True | Text | +--------------------------------------+--------+---------+-----------+ | :class:`io.BytesIO` | True | True | Binary | +--------------------------------------+--------+---------+-----------+ | :class:`io.TextIOWrapper` | True | True | Text | +--------------------------------------+--------+---------+-----------+ | :class:`io.BufferedReader` | True | False | Binary | +--------------------------------------+--------+---------+-----------+ | :class:`io.BufferedWriter` | False | True | Binary | +--------------------------------------+--------+---------+-----------+ | :class:`io.BufferedRandom` | True | True | Binary | +--------------------------------------+--------+---------+-----------+ | :func:`tempfile.TemporaryFile` | True | True | Binary | +--------------------------------------+--------+---------+-----------+ | :func:`tempfile.NamedTemporaryFile` | True | True | Binary | +--------------------------------------+--------+---------+-----------+ .. note:: When reading a list of unicode (str) lines, the input for `newline` is used to determine the number of lines in the resulting file handle, not the number of elements in the list. This is to allow composition with ``file.readlines()``. Parameters ---------- file : filepath, url, filehandle, list The input to convert to a filehandle. mode : {'r', 'w'}, optional Whether to return a readable or writable file. Conversely, this does not imply that the returned file will be unwritable or unreadable. To get a binary filehandle set `encoding` to binary. encoding : str, optional The encoding scheme to use for the file. If set to 'binary', no bytes will be translated. Otherwise this matches the behavior of :func:`io.open`. errors : str, optional Specifies how encoding and decoding errors are to be handled. This has no effect when `encoding` is binary (as there can be no errors). Otherwise this matches the behavior of :func:`io.open`. newline : {None, "", '\\n', '\\r\\n', '\\r'}, optional Matches the behavior of :func:`io.open`. compression : {'auto', 'gzip', 'bz2', None}, optional Will compress or decompress `file` depending on `mode`. If 'auto' then determining the compression of the file will be attempted and the result will be transparently decompressed. 'auto' will do nothing when writing. Other legal values will use their respective compression schemes. `compression` cannot be used with a text source. compresslevel : int (0-9 inclusive), optional The level of compression to use, will be passed to the appropriate compression handler. This is only used when writing. Returns ------- filehandle : io.TextIOBase or io.BufferedReader/Writer When `encoding='binary'` an :class:`io.BufferedReader` or :class:`io.BufferedWriter` will be returned depending on `mode`. Otherwise an implementation of :class:`io.TextIOBase` will be returned. .. note:: Any underlying resources needed to create `filehandle` are managed transparently. If `file` was closeable, garbage collection of `filehandle` will not close `file`. Calling `close` on `filehandle` will close `file`. Conversely calling `close` on `file` will cause `filehandle` to reflect a closed state. **This does not mean that a `flush` has occured for `filehandle`, there may still have been data in its buffer! Additionally, resources may not have been cleaned up properly, so ALWAYS call `close` on `filehandle` and NOT on `file`.** """ arguments = locals().copy() del arguments['file'] file, _, is_binary_file = _resolve(file, **arguments) return _munge_file(file, is_binary_file, arguments) def _munge_file(file, is_binary_file, arguments): mode = arguments.get('mode', _d['mode']) encoding = arguments.get('encoding', _d['encoding']) errors = arguments.get('errors', _d['errors']) newline = arguments.get('newline', _d['newline']) compression = arguments.get('compression', _d['compression']) is_output_binary = encoding == 'binary' newfile = file compression_handler = get_compression_handler(compression) if is_output_binary and newline is not _d['newline']: raise ValueError("Cannot use `newline` with binary encoding.") if compression is not None and not compression_handler: raise ValueError("Unsupported compression: %r" % compression) if is_binary_file: if compression: c = compression_handler(newfile, arguments) if mode == 'w': newfile = CompressedBufferedWriter(file, c.get_writer(), streamable=c.streamable) else: newfile = CompressedBufferedReader(file, c.get_reader()) if not is_output_binary: newfile = SaneTextIOWrapper(newfile, encoding=encoding, errors=errors, newline=newline) else: if compression is not None and compression != 'auto': raise ValueError("Cannot use compression with that source.") if is_output_binary: raise ValueError("Source is not a binary source") return newfile @contextmanager def _resolve_file(file, **kwargs): file, source, is_binary_file = _resolve(file, **kwargs) try: yield file, source, is_binary_file finally: if source.closeable: file.close() @contextmanager @stable(as_of="0.4.0") def open_file(file, **kwargs): r"""Context manager for :func:`skbio.io.util.open`. The signature matches :func:`open`. This context manager will not close filehandles that it did not create itself. Examples -------- Here our input isn't a filehandle and so `f` will get closed. >>> with open_file(['a\n']) as f: ... f.read() ... 'a\n' >>> f.closed True Here we provide an open file and so `f` will not get closed and neither will `file`. >>> file = io.BytesIO(b'BZh91AY&SY\x03\x89\x0c\xa6\x00\x00\x01\xc1\x00\x00' ... b'\x108\x00 \x00!\x9ah3M\x1c\xb7\x8b\xb9"\x9c(H\x01' ... b'\xc4\x86S\x00') >>> with open_file(file) as f: ... f.read() ... 'a\nb\nc\n' >>> f.closed False >>> file.closed False """ with _resolve_file(file, **kwargs) as (file, source, is_binary_file): newfile = _munge_file(file, is_binary_file, source.options) try: yield newfile finally: # As soon as we leave the above context manager file will be closed # It is important to realize that because we are closing an inner # buffer, the outer buffer will reflect that state, but it won't # get flushed as the inner buffer is oblivious to the outer # buffer's existence. if not newfile.closed: newfile.flush() _flush_compressor(newfile) def _flush_compressor(file): if isinstance(file, io.TextIOBase) and hasattr(file, 'buffer'): file = file.buffer if isinstance(file, CompressedBufferedWriter) and not file.streamable: # Some formats like BZ2 compress the entire file, and so they will # only flush once they have been closed. These kinds of files do not # close their underlying buffer, but only testing can prove that... file.raw.close() @contextmanager @stable(as_of="0.4.0") def open_files(files, **kwargs): """A plural form of :func:`open_file`.""" with ExitStack() as stack: yield [stack.enter_context(open_file(f, **kwargs)) for f in files]