Limnoria/src/utils/file.py

###
# Copyright (c) 2002-2005, Jeremiah Fincher
# Copyright (c) 2008, James McCoy
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

import os
import time
import random
import shutil
import os.path

from . import crypt
from .iter import ifilter

def contents(filename):
    with open(filename) as fd:
        return fd.read()

def open_mkdir(filename, mode='wb', *args, **kwargs):
    """filename -> file object.

    Returns a file object for filename, creating as many directories as may be
    necessary.  I.e., if the filename is ./foo/bar/baz, and . exists, and ./foo
    exists, but ./foo/bar does not exist, bar will be created before opening
    baz in it.
    """
    if mode not in ('w', 'wb'):
        raise ValueError('utils.file.open expects to write.')
    (dirname, basename) = os.path.split(filename)
    os.makedirs(dirname)
    return open(filename, mode, *args, **kwargs)

def copy(src, dst):
    """src, dst -> None

    Copies src to dst, using this module's 'open' function to open dst.
    """
    srcfd = open(src)
    dstfd = open_mkdir(dst, 'wb')
    shutil.copyfileobj(srcfd, dstfd)
    srcfd.close()
    dstfd.close()

def writeLine(fd, line):
    fd.write(line)
    if not line.endswith('\n'):
        fd.write('\n')

def readLines(filename):
    fd = open(filename)
    try:
        return [line.rstrip('\r\n') for line in fd.readlines()]
    finally:
        fd.close()

def touch(filename):
    fd = open(filename, 'w')
    fd.close()

def mktemp(suffix=''):
    """Gives a decent random string, suitable for a filename."""
    r = random.Random()
    m = crypt.md5(suffix.encode('utf8'))
    r.seed(time.time())
    s = str(r.getstate())
    period = random.random()
    now = start = time.time()
    while start + period < now:
        time.sleep() # Induce a context switch, if possible.
        now = time.time()
        m.update(str(random.random()))
        m.update(s)
        m.update(str(now))
        s = m.hexdigest()
    return crypt.sha((s + str(time.time())).encode('utf8')).hexdigest()+suffix

def nonCommentLines(fd):
    for line in fd:
        if not line.startswith('#'):
            yield line

def nonEmptyLines(fd):
    return filter(str.strip, fd)

def nonCommentNonEmptyLines(fd):
    return nonEmptyLines(nonCommentLines(fd))

def chunks(fd, size):
    return iter(lambda : fd.read(size), '')
##     chunk = fd.read(size)
##     while chunk:
##         yield chunk
##         chunk = fd.read(size)

class AtomicFile(object):
    """Used for files that need to be atomically written -- i.e., if there's a
    failure, the original file remains, unmodified.  mode must be 'w' or 'wb'"""
    class default(object): # Holder for values.
        # Callables?
        tmpDir = None
        backupDir = None
        makeBackupIfSmaller = True
        allowEmptyOverwrite = True
    def __init__(self, filename, mode='w', allowEmptyOverwrite=None,
                 makeBackupIfSmaller=None, tmpDir=None, backupDir=None):
        if tmpDir is None:
            tmpDir = force(self.default.tmpDir)
        if backupDir is None:
            backupDir = force(self.default.backupDir)
        if makeBackupIfSmaller is None:
            makeBackupIfSmaller = force(self.default.makeBackupIfSmaller)
        if allowEmptyOverwrite is None:
            allowEmptyOverwrite = force(self.default.allowEmptyOverwrite)
        if mode not in ('w', 'wb'):
            raise ValueError(format('Invalid mode: %q', mode))
        self.rolledback = False
        self.allowEmptyOverwrite = allowEmptyOverwrite
        self.makeBackupIfSmaller = makeBackupIfSmaller
        self.filename = filename
        self.backupDir = backupDir
        if tmpDir is None:
            # If not given a tmpDir, we'll just put a random token on the end
            # of our filename and put it in the same directory.
            self.tempFilename = '%s.%s' % (self.filename, mktemp())
        else:
            # If given a tmpDir, we'll get the basename (just the filename, no
            # directory), put our random token on the end, and put it in tmpDir
            tempFilename = '%s.%s' % (os.path.basename(self.filename), mktemp())
            self.tempFilename = os.path.join(tmpDir, tempFilename)
        # This doesn't work because of the uncollectable garbage effect.
        # self.__parent = super(AtomicFile, self)
        self._fd = open(self.tempFilename, mode)

    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            self.rollback()
        else:
            self.close()


    @property
    def closed(self):
        return self._fd.closed

    def write(self, data):
        return self._fd.write(data)

    def writelines(self, lines):
        return self._fd.writelines(lines)

    def rollback(self):
        if not self.closed:
            self._fd.close()
            if os.path.exists(self.tempFilename):
                os.remove(self.tempFilename)
            self.rolledback = True

    def seek(self, offset):
        return self._fd.seek(offset)

    def tell(self):
        return self._fd.tell()

    def flush(self):
        return self._fd.flush()

    def close(self):
        if not self.rolledback:
            self._fd.close()
            # We don't mind writing an empty file if the file we're overwriting
            # doesn't exist.
            newSize = os.path.getsize(self.tempFilename)
            originalExists = os.path.exists(self.filename)
            if newSize or self.allowEmptyOverwrite or not originalExists:
                if originalExists:
                    oldSize = os.path.getsize(self.filename)
                    if self.makeBackupIfSmaller and newSize < oldSize and \
                            self.backupDir != '/dev/null':
                        now = int(time.time())
                        backupFilename = '%s.backup.%s' % (self.filename, now)
                        if self.backupDir is not None:
                            backupFilename = os.path.basename(backupFilename)
                            backupFilename = os.path.join(self.backupDir,
                                                          backupFilename)
                        shutil.copy(self.filename, backupFilename)
                # We use shutil.move here instead of os.rename because
                # the latter doesn't work on Windows when self.filename
                # (the target) already exists.  shutil.move handles those
                # intricacies for us.

                # This raises IOError if we can't write to the file.  Since
                # in *nix, it only takes write perms to the *directory* to
                # rename a file (and shutil.move will use os.rename if
                # possible), we first check if we have the write permission
                # and only then do we write.
                fd = open(self.filename, 'a')
                fd.close()
                shutil.move(self.tempFilename, self.filename)

        else:
            raise ValueError('AtomicFile.close called after rollback.')

    def __del__(self):
        # We rollback because if we're deleted without being explicitly closed,
        # that's bad.  We really should log this here, but as of yet we've got
        # no logging facility in utils.  I've got some ideas for this, though.
        self.rollback()

# vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79: