~launchpad-pqm/launchpad/devel

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
# Copyright 2009-2010 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

"""Utilities for doing the sort of thing the os module does."""

__metaclass__ = type
__all__ = [
    'ensure_directory_exists',
    'kill_by_pidfile',
    'open_for_writing',
    'override_environ',
    'remove_if_exists',
    'remove_tree',
    'two_stage_kill',
    'until_no_eintr',
    ]

from contextlib import contextmanager
import errno
import os.path
import shutil
import socket

from canonical.launchpad.daemons.tachandler import (
    kill_by_pidfile,
    remove_if_exists,
    two_stage_kill,
    )


def remove_tree(path):
    """Remove the tree at 'path' from disk."""
    if os.path.exists(path):
        shutil.rmtree(path)


def set_environ(new_values):
    """Set the environment variables as specified by new_values.

    :return: a dict of the old values
    """
    old_values = {}
    for name, value in new_values.iteritems():
        old_values[name] = os.environ.get(name)
        if value is None:
            if old_values[name] is not None:
                del os.environ[name]
        else:
            os.environ[name] = value
    return old_values


@contextmanager
def override_environ(**kwargs):
    """Override environment variables with the kwarg values.

    If a value is None, the environment variable is deleted.  Variables are
    restored to their previous state when exiting the context.
    """
    old_values = set_environ(kwargs)
    try:
        yield
    finally:
        set_environ(old_values)


def until_no_eintr(retries, function, *args, **kwargs):
    """Run 'function' until it doesn't raise EINTR errors.

    :param retries: The maximum number of times to try running 'function'.
    :param function: The function to run.
    :param *args: Arguments passed to the function.
    :param **kwargs: Keyword arguments passed to the function.
    :return: The return value of 'function'.
    """
    if not retries:
        return
    for i in range(retries):
        try:
            return function(*args, **kwargs)
        except (IOError, OSError), e:
            if e.errno == errno.EINTR:
                continue
            raise
        except socket.error, e:
            # In Python 2.6 we can use IOError instead.  It also has
            # reason.errno but we might be using 2.5 here so use the
            # index hack.
            if e[0] == errno.EINTR:
                continue
            raise
    else:
        raise


def ensure_directory_exists(directory, mode=0777):
    """Create 'directory' if it doesn't exist.

    :return: True if the directory had to be created, False otherwise.
    """
    try:
        os.makedirs(directory, mode=mode)
    except OSError, e:
        if e.errno == errno.EEXIST:
            return False
        raise
    return True


def open_for_writing(filename, mode, dirmode=0777):
    """Open 'filename' for writing, creating directories if necessary.

    :param filename: The path of the file to open.
    :param mode: The mode to open the filename with. Should be 'w', 'a' or
        something similar. See ``open`` for more details. If you pass in
        a read-only mode (e.g. 'r'), then we'll just accept that and return
        a read-only file-like object.
    :param dirmode: The mode to use to create directories, if necessary.
    :return: A file-like object that can be used to write to 'filename'.
    """
    try:
        return open(filename, mode)
    except IOError, e:
        if e.errno == errno.ENOENT:
            os.makedirs(os.path.dirname(filename), mode=dirmode)
            return open(filename, mode)