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
|
"""
State
-----
Classes for storing a program's state.
"""
import os
import atexit
import yaml
try:
from sqlalchemy import create_engine, Column, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, reconstructor
Base = declarative_base()
HAVE_SQL_ALCHEMY = True
except:
HAVE_SQL_ALCHEMY = False
class State(object):
"""
A program's state.
Contains a dictionary that can be periodically saved and restored at startup.
Maybe later this will be subclassed with database connectors and whatnot,
but for now it'll just save to a yaml file.
"""
@classmethod
def state(cls, *args, **kwargs):
return cls(*args, **kwargs)
def __init__(self, path=None):
self.path = path
self.d = {}
self.load()
atexit.register(self._exit_handler)
def __enter__(self):
self.load()
return self
def __exit__(self, type, value, traceback):
self.save()
def __getitem__(self, key):
try:
return self.d[key]
except KeyError:
return None
def __setitem__(self, key, value):
self.d[key] = value
def _exit_handler(self):
self.save()
def save(self):
"""
Save the state to a file.
"""
with open(self.path, 'w') as f:
f.write(yaml.dump(dict(self.d)))
def load(self):
"""
Load a saved state file.
"""
if os.path.exists(self.path):
with open(self.path, 'r') as f:
d = yaml.safe_load(f.read().replace('\t', ' '*4))
# don't clobber self.d if we successfully opened the state file
# but it was empty
if d:
self.d = d
def cleanup(self):
"""
Clean up the saved state.
"""
if os.path.exists(self.path):
os.remove(self.path)
if HAVE_SQL_ALCHEMY:
class DBState(State, Base):
"""
State stored in a database, using SQLAlchemy.
"""
__tablename__ = 'state'
id = Column(Integer, primary_key=True)
data = Column(String)
session = None
@classmethod
def state(cls, url=None, *args, **kwargs):
if not cls.session:
engine = create_engine(url, echo=False)
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)
cls.session = Session()
inst = cls.session.query(DBState).first()
if inst:
return inst
else:
return cls(*args, **kwargs)
def __init__(self, *args, **kwargs):
super(DBState, self).__init__(*args, **kwargs)
self.d = {}
self.data = '{}'
def save(self):
self.data = yaml.dump(self.d)
self.session.add(self)
self.session.commit()
@reconstructor
def load(self):
if self.data:
self.d = yaml.safe_load(self.data)
def cleanup(self):
self.d = {}
self.save()
|
