util.py 2.07 KB
Newer Older
1
"""Common system utilities"""
2 3 4 5
from __future__ import absolute_import as _abs
import os
import tempfile
import shutil
6 7 8 9 10
try:
    import fcntl
except ImportError:
    fcntl = None

11 12

class TempDirectory(object):
13 14 15 16
    """Helper object to manage temp directory during testing.

    Automatically removes the directory when it went out of scope.
    """
17 18
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp()
Tianqi Chen committed
19
        self._rmtree = shutil.rmtree
20

21 22 23
    def remove(self):
        """Remote the tmp dir"""
        if self.temp_dir:
Hu Shiwen committed
24
            self._rmtree(self.temp_dir, ignore_errors=True)
25 26
            self.temp_dir = None

27
    def __del__(self):
28
        self.remove()
29 30 31 32 33 34 35 36

    def relpath(self, name):
        """Relative path in temp dir

        Parameters
        ----------
        name : str
            The name of the file.
37 38 39 40 41

        Returns
        -------
        path : str
            The concatenated path.
42 43 44
        """
        return os.path.join(self.temp_dir, name)

45
    def listdir(self):
46
        """List contents in the dir.
47 48 49 50 51 52

        Returns
        -------
        names : list
            The content of directory
        """
53
        return os.listdir(self.temp_dir)
54

55

56
def tempdir():
57
    """Create temp dir which deletes the contents when exit.
58 59 60 61 62 63 64

    Returns
    -------
    temp : TempDirectory
        The temp directory object
    """
    return TempDirectory()
65 66 67 68 69 70 71 72 73 74 75 76


class FileLock(object):
    """File lock object

    Parameters
    ----------
    path : str
        The path to the lock
    """
    def __init__(self, path):
        self.lock_file = open(path, "w")
77 78
        if fcntl:
            fcntl.lockf(self.lock_file, fcntl.LOCK_EX)
79 80 81 82 83


    def release(self):
        """Release the lock"""
        if self.lock_file:
84 85
            if fcntl:
                fcntl.lockf(self.lock_file, fcntl.LOCK_UN)
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
            self.lock_file.close()
            self.lock_file = None

def filelock(path):
    """Create a file lock which locks on path

    Parameters
    ----------
    path : str
        The path to the lock

    Returns
    -------
    lock : File lock object
    """
    return FileLock(path)