# -*- test-case-name: twisted.test.test_sob -*- # Copyright (c) Twisted Matrix Laboratories. # See LICENSE for details. # """ Save and load Small OBjects to and from files, using various formats. Maintainer: Moshe Zadka """ from __future__ import division, absolute_import import os import sys try: import cPickle as pickle except ImportError: import pickle from twisted.python import log, runtime from twisted.persisted import styles from zope.interface import implementer, Interface class IPersistable(Interface): """An object which can be saved in several formats to a file""" def setStyle(style): """Set desired format. @type style: string (one of 'pickle' or 'source') """ def save(tag=None, filename=None, passphrase=None): """Save object to file. @type tag: string @type filename: string @type passphrase: string """ @implementer(IPersistable) class Persistent: style = "pickle" def __init__(self, original, name): self.original = original self.name = name def setStyle(self, style): """Set desired format. @type style: string (one of 'pickle' or 'source') """ self.style = style def _getFilename(self, filename, ext, tag): if filename: finalname = filename filename = finalname + "-2" elif tag: filename = "%s-%s-2.%s" % (self.name, tag, ext) finalname = "%s-%s.%s" % (self.name, tag, ext) else: filename = "%s-2.%s" % (self.name, ext) finalname = "%s.%s" % (self.name, ext) return finalname, filename def _saveTemp(self, filename, dumpFunc): with open(filename, 'wb') as f: dumpFunc(self.original, f) def _getStyle(self): if self.style == "source": from twisted.persisted.aot import jellyToSource as dumpFunc ext = "tas" else: def dumpFunc(obj, file): pickle.dump(obj, file, 2) ext = "tap" return ext, dumpFunc def save(self, tag=None, filename=None, passphrase=None): """Save object to file. @type tag: string @type filename: string @type passphrase: string """ ext, dumpFunc = self._getStyle() if passphrase is not None: raise TypeError("passphrase must be None") ext = 'e' + ext finalname, filename = self._getFilename(filename, ext, tag) log.msg("Saving "+self.name+" application to "+finalname+"...") self._saveTemp(filename, dumpFunc) if runtime.platformType == "win32" and os.path.isfile(finalname): os.remove(finalname) os.rename(filename, finalname) log.msg("Saved.") # "Persistant" has been present since 1.0.7, so retain it for compatibility Persistant = Persistent class _EverythingEphemeral(styles.Ephemeral): initRun = 0 def __init__(self, mainMod): """ @param mainMod: The '__main__' module that this class will proxy. """ self.mainMod = mainMod def __getattr__(self, key): try: return getattr(self.mainMod, key) except AttributeError: if self.initRun: raise else: log.msg("Warning! Loading from __main__: %s" % key) return styles.Ephemeral() def load(filename, style): """Load an object from a file. Deserialize an object from a file. The file can be encrypted. @param filename: string @param style: string (one of 'pickle' or 'source') """ mode = 'r' if style=='source': from twisted.persisted.aot import unjellyFromSource as _load else: _load, mode = pickle.load, 'rb' fp = open(filename, mode) ee = _EverythingEphemeral(sys.modules['__main__']) sys.modules['__main__'] = ee ee.initRun = 1 with fp: try: value = _load(fp) finally: # restore __main__ if an exception is raised. sys.modules['__main__'] = ee.mainMod styles.doUpgrade() ee.initRun = 0 persistable = IPersistable(value, None) if persistable is not None: persistable.setStyle(style) return value def loadValueFromFile(filename, variable): """Load the value of a variable in a Python file. Run the contents of the file in a namespace and return the result of the variable named C{variable}. @param filename: string @param variable: string """ with open(filename, 'r') as fileObj: data = fileObj.read() d = {'__file__': filename} codeObj = compile(data, filename, "exec") eval(codeObj, d, d) value = d[variable] return value def guessType(filename): ext = os.path.splitext(filename)[1] return { '.tac': 'python', '.etac': 'python', '.py': 'python', '.tap': 'pickle', '.etap': 'pickle', '.tas': 'source', '.etas': 'source', }[ext] __all__ = ['loadValueFromFile', 'load', 'Persistent', 'Persistant', 'IPersistable', 'guessType']