imcflibs.iotools

  1"""I/O related functions."""
  2
  3import glob
  4import zipfile
  5
  6import os
  7from os.path import splitext, join
  8
  9from .log import LOG as log
 10from .strtools import flatten
 11
 12
 13def filehandle(fname, mode="r"):
 14    """Make sure a variable is either a filehandle or create one from it.
 15
 16    This function takes a variable and checks whether it is already a
 17    filehandle with the desired mode or a string that can be turned into a
 18    filehandle with that mode. This can be used e.g. to make functions agnostic
 19    against being supplied a file-type parameter that was gathered via argparse
 20    (then it's already a filehandle) or as a plain string.
 21
 22    Parameters
 23    ----------
 24    fname : str or filehandle
 25    mode : str
 26        The desired mode of the filehandle (default=read).
 27
 28    Returns
 29    -------
 30    A valid (open) filehandle with the given mode. Raises an IOError
 31    otherwise.
 32
 33    Example
 34    -------
 35    >>> fname = __file__
 36    >>> type(fname)
 37    <type 'str'>
 38    >>> type(filehandle(fname))
 39    <type 'file'>
 40    >>> fh = open(__file__, 'r')
 41    >>> type(fh)
 42    <type 'file'>
 43    >>> type(filehandle(fh))
 44    <type 'file'>
 45    """
 46    log.debug(type(fname))
 47    if type(fname).__name__ == "str":
 48        try:
 49            return open(fname, mode)
 50        except IOError as err:
 51            message = "can't open '%s': %s"
 52            raise SystemExit(message % (fname, err))
 53    elif type(fname).__name__ == "file":
 54        if fname.mode != mode:
 55            message = "mode mismatch: %s != %s"
 56            raise IOError(message % (fname.mode, mode))
 57        return fname
 58    else:
 59        message = "unknown data type (expected string or filehandle): %s"
 60        raise SystemExit(message % type(fname))
 61
 62
 63def readtxt(fname, path="", flat=False):
 64    """Commodity function for reading text files plain or zipped.
 65
 66    Read a text file line by line either plainly from a directory or a .zip or
 67    .jar file. Return as a list of strings or optionally flattened into a
 68    single string.
 69
 70    BEWARE: this is NOT intended for HUGE text files as it actually reads them
 71    in and returns the content, not a handle to the reader itself!
 72
 73    Parameters
 74    ----------
 75    fname : str
 76        The name of the file to read in. Can be a full or relative path if
 77        desired. For automatic archive handling use the 'path' parameter.
 78    path : str (optional)
 79        The directory where to look for the file. If the string has the suffix
 80        '.zip' or '.jar' an archive is assumed and the corresponding mechanisms
 81        are used to read 'fname' from within this archive.
 82    flat : bool (optional)
 83        Used to request a flattened string instead of a list of strings.
 84
 85    Returns
 86    -------
 87    txt : str or list(str)
 88
 89    Example
 90    -------
 91    >>> readtxt('foo', '/tmp/archive.zip', flat=True)
 92    ... # doctest: +SKIP
 93    """
 94    zipread = None
 95    suffix = splitext(path)[1].lower()
 96    if (suffix == ".zip") or (suffix == ".jar"):
 97        # ZipFile only works as a context manager from Python 2.7 on
 98        # tag:python25
 99        zipread = zipfile.ZipFile(path, "r")
100        fin = zipread.open(fname)
101    else:
102        fin = open(join(path, fname), "r")
103    txt = fin.readlines()  # returns file as a list, one entry per line
104    if flat:
105        txt = flatten(txt)
106    fin.close()
107    if zipread is not None:
108        zipread.close()
109    return txt