GitLab is used only for code review, issue tracking and project management. Canonical locations for source code are still https://gitweb.torproject.org/ https://git.torproject.org/ and git-rw.torproject.org.

state.py 3.73 KB
Newer Older
1 2 3 4 5 6
from sbws.util.filelock import FileLock
import os
import json


class State:
Matt Traudt's avatar
Matt Traudt committed
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
    '''
    State allows one to atomically access and update a simple state file on
    disk across threads and across processes.

    To put it blunty, due to limited developer time and his inability to
    quickly find a way to safely access and update more complex data types
    (namely, collections like list, set, and dict), you may only store simple
    types of data as enumerated in _ALLOWED_TYPES. Keys must be strings.

    Data is stored as JSON on disk in the provided file file.

    >>> state = State('foo.state')
    >>> # state == {}

    >>> state['linux'] = True
    >>> # 'foo.state' now exists on disk with the JSON for {'linux': True}

    >>> # We read 'foo.state' from disk in order to get the most up-to-date
    >>> #     state info. Pretend another process has updated 'linux' to be
    >>> #     False
    >>> state['linux']
    >>> # returns False

    >>> # Pretend another process has added the user's age to the state file.
    >>> #     As before, we read the state file from disk for the most
    >>> #     up-to-date info.
    >>> state['age']
    >>> # Returns 14

    >>> # We now set their name. We read the state file first, set the option,
    >>> #     and then write it out.
    >>> state['name'] = 'John'

    >>> # We can do many of the same things with a State object as with a dict
    >>> for key in state: print(key)
    >>> # Prints 'linux', 'age', and 'name'
    '''
44 45 46 47 48 49 50
    _ALLOWED_TYPES = (int, float, str, bool, type(None))

    def __init__(self, fname):
        self._fname = fname
        self._state = self._read()

    def _read(self):
51 52
        if not os.path.exists(self._fname):
            return {}
53 54 55 56 57 58 59
        with FileLock(self._fname):
            with open(self._fname, 'rt') as fd:
                return json.load(fd)

    def _write(self):
        with FileLock(self._fname):
            with open(self._fname, 'wt') as fd:
60
                return json.dump(self._state, fd, indent=4)
61 62 63 64 65

    def __len__(self):
        self._state = self._read()
        return self._state.__len__()

66 67 68 69 70
    def get(self, key, d=None):
        """
        Implements a dictionary ``get`` method reading and locking
        a json file.
        """
juga  's avatar
juga committed
71 72 73 74
        if not isinstance(key, str):
            raise TypeError(
                'Keys must be strings. %s is a %s' % (key, type(key)))
        self._state = self._read()
75
        return self._state.get(key, d)
juga  's avatar
juga committed
76

77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
    def __getitem__(self, key):
        if not isinstance(key, str):
            raise TypeError(
                'Keys must be strings. %s is a %s' % (key, type(key)))
        self._state = self._read()
        return self._state.__getitem__(key)

    def __delitem__(self, key):
        if not isinstance(key, str):
            raise TypeError(
                'Keys must be strings. %s is a %s' % (key, type(key)))
        self._state = self._read()
        self._state.__delitem__(key)
        self._write()

    def __setitem__(self, key, value):
        if not isinstance(key, str):
            raise TypeError(
                'Keys must be strings. %s is a %s' % (key, type(key)))
        if type(value) not in State._ALLOWED_TYPES:
            raise TypeError(
                'May only store value with type in %s, not %s' %
                (State._ALLOWED_TYPES, type(value)))
100 101 102
        # NOTE: important, read the file before setting the key,
        # otherwise if other instances are creating other keys, they're lost.
        self._state = self._read()
103 104 105 106 107 108 109 110 111 112
        self._state.__setitem__(key, value)
        self._write()

    def __iter__(self):
        self._state = self._read()
        return self._state.__iter__()

    def __contains__(self, item):
        self._state = self._read()
        return self._state.__contains__(item)