aboutsummaryrefslogtreecommitdiff
blob: fdd0f23e6069a776fbfbfca0f6624ca09e39d8dd (plain)
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
"""
file related operations, mainly reading
"""

import mmap
import os
from functools import partial

from . import _fileutils, data_source
from .compatibility import IGNORED_EXCEPTIONS
from .currying import pretty_docs
from .klass import GetAttrProxy


def touch(fname: str, mode: int = 0o644, dir_fd=None, **kwargs):
    """touch(1) equivalent

    :param fname: file path
    :param mode: file mode

    See os.utime for other supported arguments.
    """
    flags = os.O_CREAT | os.O_APPEND

    with os.fdopen(os.open(fname, flags, mode, dir_fd=dir_fd)) as f:
        os.utime(
            f.fileno() if os.utime in os.supports_fd else fname,
            dir_fd=None if os.supports_fd else dir_fd,
            **kwargs,
        )


def mmap_or_open_for_read(path: str):
    size = os.stat(path).st_size
    if size == 0:
        return (None, data_source.bytes_ro_StringIO(b""))
    fd = None
    try:
        fd = os.open(path, os.O_RDONLY)
        return (
            _fileutils.mmap_and_close(fd, size, mmap.MAP_SHARED, mmap.PROT_READ),
            None,
        )
    except IGNORED_EXCEPTIONS:
        raise
    except:
        try:
            os.close(fd)
        except EnvironmentError:
            pass
        raise


class AtomicWriteFile_mixin:
    """File class that stores the changes in a tempfile.

    Upon invocation of the close method, this class will use
    :py:func:`os.rename` to atomically replace the destination.

    Similar to file protocol behavior, except that close *must*
    be called for the changes to be made live,

    If along the way it's decided that these changes should be discarded,
    invoke :py:func:`AtomicWriteFile.discard`; this will close the file
    without updating the target.

    If this object falls out of memory without ever being discarded nor
    closed, the contents are discarded and a warning is issued.
    """

    def __init__(self, fp, binary=False, perms=None, uid=-1, gid=-1):
        """
        :param fp: filepath to write to upon close
        :param binary: should we open the file in binary mode?
        :param perms: if specified, permissions we should force for the file.
        :param uid: if specified, the uid to force for the file.
        :param gid: if specified, the uid to force for the file.
        """
        self._is_finalized = True
        if binary:
            file_mode = "wb"
        else:
            file_mode = "w"
        self._computed_mode = file_mode
        fp = os.path.realpath(fp)
        self._original_fp = fp
        self._temp_fp = os.path.join(
            os.path.dirname(fp), ".update." + os.path.basename(fp)
        )
        old_umask = None
        if perms:
            # give it just write perms
            old_umask = os.umask(0o0200)
        try:
            self._actual_init()
        finally:
            if old_umask is not None:
                os.umask(old_umask)
        self._is_finalized = False
        if perms:
            os.chmod(self._temp_fp, perms)
        if (gid, uid) != (-1, -1):
            os.chown(self._temp_fp, uid, gid)

    def discard(self):
        """If we've not already flushed our changes to the target, discard them
        and close this file handle."""
        if not self._is_finalized:
            self._real_close()
            os.unlink(self._temp_fp)
            self._is_finalized = True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        if exc is not None:
            self.discard()
        else:
            self.close()

    def close(self):
        """Close this file handle, atomically updating the target in the process.

        Note that if we're already closed, this method does nothing
        """
        if not self._is_finalized:
            self._real_close()
            os.rename(self._temp_fp, self._original_fp)
            self._is_finalized = True

    def __del__(self):
        self.discard()


class AtomicWriteFile(AtomicWriteFile_mixin):
    __doc__ = AtomicWriteFile_mixin.__doc__

    def _actual_init(self):
        self.raw = open(self._temp_fp, mode=self._computed_mode)

    def _real_close(self):
        if hasattr(self, "raw"):
            return self.raw.close()
        return None

    __getattr__ = GetAttrProxy("raw")


def _mk_pretty_derived_func(func, name_base: str, name: str, *args, **kwds):
    if name:
        name = "_" + name
    return pretty_docs(partial(func, *args, **kwds), name=name_base + name)


_mk_readfile = partial(_mk_pretty_derived_func, _fileutils.native_readfile, "readfile")

readfile_ascii = _mk_readfile("ascii", "rt")
readfile_bytes = _mk_readfile("bytes", "rb")
readfile_utf8 = _mk_readfile("utf8", "r", encoding="utf8")
readfile = readfile_utf8


_mk_readlines = partial(
    _mk_pretty_derived_func, _fileutils.native_readlines, "readlines"
)

readlines_ascii = _mk_readlines("ascii", "r", encoding="ascii")
readlines_bytes = _mk_readlines("bytes", "rb")
readlines_utf8 = _mk_readlines("utf8", "r", encoding="utf8")
readlines = readlines_utf8