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
|
"""
PC-BASIC - state.py
Support for pickling emulator state
(c) 2014--2020 Rob Hagemans
This file is released under the GNU GPL version 3 or later.
"""
try:
import cPickle as pickle
except ImportError:
import pickle
import os
import io
import sys
import zlib
import struct
import codecs
import logging
from contextlib import contextmanager
from .basic import VERSION
from .compat import PY2, copyreg, stdio
# session file header
HEADER_FORMAT = '<LIIIII'
HEADER_KEYS = [
'checksum', 'format_version', 'python_major', 'python_minor', 'pcbasic_major', 'pcbasic_minor'
]
HEADER = {
# increment this if we change the format of the session file
'format_version': 2,
'python_major': sys.version_info.major,
'python_minor': sys.version_info.minor,
'pcbasic_major': int(VERSION.split(u'.')[0]),
'pcbasic_minor': int(VERSION.split(u'.')[1]),
}
@contextmanager
def manage_state(session, state_file, do_resume):
"""Resume a session if requested; save upon exit"""
if do_resume and state_file:
try:
session = load_session(state_file).attach(session.interface)
except Exception as e:
# if we were told to resume but can't, give up
logging.fatal('Failed to resume session from %s: %s', state_file, e)
sys.exit(1)
try:
yield session
finally:
if state_file:
try:
save_session(session, state_file)
except Exception as e:
logging.error('Failed to save session to %s: %s', state_file, e)
def unpickle_bytesio(value, pos):
"""Unpickle a file object."""
stream = io.BytesIO(value)
stream.seek(pos)
return stream
def pickle_bytesio(f):
"""Pickle a BytesIO object."""
return unpickle_bytesio, (f.getvalue(), f.tell())
def unpickle_file(name, mode, pos):
"""Unpickle a file object."""
if name is None:
if mode == 'rb' or (PY2 and mode == 'r'):
return stdio.stdin.buffer
elif mode == 'r':
return stdio.stdin
elif mode == 'wb' or (PY2 and mode == 'w'):
return stdio.stdout.buffer
elif mode == 'w':
return stdio.stdout
try:
if 'w' in mode and pos > 0:
# preserve existing contents of writable file
with io.open(name, 'rb') as f:
buf = f.read(pos)
f = io.open(name, mode)
f.write(buf)
else:
f = io.open(name, mode)
if pos > 0:
f.seek(pos)
except IOError:
pass
else:
return f
logging.warning('Could not re-open file %s. Replacing with null file.', name)
return io.open(os.devnull, mode)
def pickle_file(f):
"""Pickle a file object."""
if f in (
sys.stdout, sys.stdin,
stdio.stdout, stdio.stdin,
stdio.stdout.buffer, stdio.stdin.buffer
):
return unpickle_file, (None, f.mode, -1)
try:
return unpickle_file, (f.name, f.mode, f.tell())
except (IOError, ValueError):
# IOError: not seekable
# ValueError: closed
return unpickle_file, (f.name, f.mode, -1)
# register the picklers for file and cStringIO
if PY2:
copyreg.pickle(file, pickle_file) # pylint: disable=undefined-variable
copyreg.pickle(io.BufferedReader, pickle_file)
copyreg.pickle(io.BufferedWriter, pickle_file)
copyreg.pickle(io.TextIOWrapper, pickle_file)
copyreg.pickle(io.BufferedRandom, pickle_file)
copyreg.pickle(io.BytesIO, pickle_bytesio)
# patch codecs.StreamReader and -Writer
if PY2:
def patched_getstate(self):
return vars(self)
def patched_setstate(self, dict):
vars(self).update(dict)
for streamclass in (codecs.StreamReader, codecs.StreamWriter, codecs.StreamReaderWriter):
streamclass.__getstate__ = patched_getstate
streamclass.__setstate__ = patched_setstate
def load_session(state_file):
"""Read state from a compressed pickle."""
with open(state_file, 'rb') as in_file:
header = in_file.read(struct.calcsize(HEADER_FORMAT))
blob = in_file.read()
# mask checksum to deal with different signs on Py2/Py3
# see https://docs.python.org/3.5/library/zlib.html#zlib.crc32
checksum = zlib.crc32(blob) & 0xffffffff
try:
header_dict = dict(zip(HEADER_KEYS, struct.unpack(HEADER_FORMAT, header)))
except struct.error:
raise ValueError('session file header corrupted')
# check blob integrity
if checksum != header_dict['checksum']:
raise ValueError('session file corrupted')
if (
HEADER['python_major'] != header_dict['python_major']
or HEADER['python_minor'] != header_dict['python_minor']
):
raise ValueError('session file stored with different Python version')
if (
HEADER['pcbasic_major'] != header_dict['pcbasic_major']
or HEADER['pcbasic_minor'] != header_dict['pcbasic_minor']
):
raise ValueError('session file stored with different PC-BASIC version')
session = pickle.loads(zlib.decompress(blob))
return session
def save_session(obj, state_file):
"""Write state to a compressed pickle."""
blob = zlib.compress(pickle.dumps(obj, pickle.HIGHEST_PROTOCOL))
checksum = zlib.crc32(blob) & 0xffffffff
header_dict = dict(checksum=checksum, **HEADER)
header = struct.pack(HEADER_FORMAT, *(header_dict[_key] for _key in HEADER_KEYS))
with open(state_file, 'wb') as out_file:
out_file.write(header)
out_file.write(blob)
|