You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

263 lines
8.0 KiB

4 years ago
  1. #!/usr/bin/env python
  2. #
  3. # Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
  4. # Copyright (c) 2008-2016 California Institute of Technology.
  5. # Copyright (c) 2016-2018 The Uncertainty Quantification Foundation.
  6. # License: 3-clause BSD. The full license text is available at:
  7. # - https://github.com/uqfoundation/dill/blob/master/LICENSE
  8. """
  9. Methods for serialized objects (or source code) stored in temporary files
  10. and file-like objects.
  11. """
  12. #XXX: better instead to have functions write to any given file-like object ?
  13. #XXX: currently, all file-like objects are created by the function...
  14. __all__ = ['dump_source', 'dump', 'dumpIO_source', 'dumpIO',\
  15. 'load_source', 'load', 'loadIO_source', 'loadIO',\
  16. 'capture']
  17. import contextlib
  18. from ._dill import PY3
  19. @contextlib.contextmanager
  20. def capture(stream='stdout'):
  21. """builds a context that temporarily replaces the given stream name
  22. >>> with capture('stdout') as out:
  23. ... print "foo!"
  24. ...
  25. >>> print out.getvalue()
  26. foo!
  27. """
  28. import sys
  29. if PY3:
  30. from io import StringIO
  31. else:
  32. from StringIO import StringIO
  33. orig = getattr(sys, stream)
  34. setattr(sys, stream, StringIO())
  35. try:
  36. yield getattr(sys, stream)
  37. finally:
  38. setattr(sys, stream, orig)
  39. def b(x): # deal with b'foo' versus 'foo'
  40. import codecs
  41. return codecs.latin_1_encode(x)[0]
  42. def load_source(file, **kwds):
  43. """load an object that was stored with dill.temp.dump_source
  44. file: filehandle
  45. alias: string name of stored object
  46. mode: mode to open the file, one of: {'r', 'rb'}
  47. >>> f = lambda x: x**2
  48. >>> pyfile = dill.temp.dump_source(f, alias='_f')
  49. >>> _f = dill.temp.load_source(pyfile)
  50. >>> _f(4)
  51. 16
  52. """
  53. alias = kwds.pop('alias', None)
  54. mode = kwds.pop('mode', 'r')
  55. fname = getattr(file, 'name', file) # fname=file.name or fname=file (if str)
  56. source = open(fname, mode=mode, **kwds).read()
  57. if not alias:
  58. tag = source.strip().splitlines()[-1].split()
  59. if tag[0] != '#NAME:':
  60. stub = source.splitlines()[0]
  61. raise IOError("unknown name for code: %s" % stub)
  62. alias = tag[-1]
  63. local = {}
  64. exec(source, local)
  65. _ = eval("%s" % alias, local)
  66. return _
  67. def dump_source(object, **kwds):
  68. """write object source to a NamedTemporaryFile (instead of dill.dump)
  69. Loads with "import" or "dill.temp.load_source". Returns the filehandle.
  70. >>> f = lambda x: x**2
  71. >>> pyfile = dill.temp.dump_source(f, alias='_f')
  72. >>> _f = dill.temp.load_source(pyfile)
  73. >>> _f(4)
  74. 16
  75. >>> f = lambda x: x**2
  76. >>> pyfile = dill.temp.dump_source(f, dir='.')
  77. >>> modulename = os.path.basename(pyfile.name).split('.py')[0]
  78. >>> exec('from %s import f as _f' % modulename)
  79. >>> _f(4)
  80. 16
  81. Optional kwds:
  82. If 'alias' is specified, the object will be renamed to the given string.
  83. If 'prefix' is specified, the file name will begin with that prefix,
  84. otherwise a default prefix is used.
  85. If 'dir' is specified, the file will be created in that directory,
  86. otherwise a default directory is used.
  87. If 'text' is specified and true, the file is opened in text
  88. mode. Else (the default) the file is opened in binary mode. On
  89. some operating systems, this makes no difference.
  90. NOTE: Keep the return value for as long as you want your file to exist !
  91. """ #XXX: write a "load_source"?
  92. from .source import importable, getname
  93. import tempfile
  94. kwds.pop('suffix', '') # this is *always* '.py'
  95. alias = kwds.pop('alias', '') #XXX: include an alias so a name is known
  96. name = str(alias) or getname(object)
  97. name = "\n#NAME: %s\n" % name
  98. #XXX: assumes kwds['dir'] is writable and on $PYTHONPATH
  99. file = tempfile.NamedTemporaryFile(suffix='.py', **kwds)
  100. file.write(b(''.join([importable(object, alias=alias),name])))
  101. file.flush()
  102. return file
  103. def load(file, **kwds):
  104. """load an object that was stored with dill.temp.dump
  105. file: filehandle
  106. mode: mode to open the file, one of: {'r', 'rb'}
  107. >>> dumpfile = dill.temp.dump([1, 2, 3, 4, 5])
  108. >>> dill.temp.load(dumpfile)
  109. [1, 2, 3, 4, 5]
  110. """
  111. import dill as pickle
  112. mode = kwds.pop('mode', 'rb')
  113. name = getattr(file, 'name', file) # name=file.name or name=file (if str)
  114. return pickle.load(open(name, mode=mode, **kwds))
  115. def dump(object, **kwds):
  116. """dill.dump of object to a NamedTemporaryFile.
  117. Loads with "dill.temp.load". Returns the filehandle.
  118. >>> dumpfile = dill.temp.dump([1, 2, 3, 4, 5])
  119. >>> dill.temp.load(dumpfile)
  120. [1, 2, 3, 4, 5]
  121. Optional kwds:
  122. If 'suffix' is specified, the file name will end with that suffix,
  123. otherwise there will be no suffix.
  124. If 'prefix' is specified, the file name will begin with that prefix,
  125. otherwise a default prefix is used.
  126. If 'dir' is specified, the file will be created in that directory,
  127. otherwise a default directory is used.
  128. If 'text' is specified and true, the file is opened in text
  129. mode. Else (the default) the file is opened in binary mode. On
  130. some operating systems, this makes no difference.
  131. NOTE: Keep the return value for as long as you want your file to exist !
  132. """
  133. import dill as pickle
  134. import tempfile
  135. file = tempfile.NamedTemporaryFile(**kwds)
  136. pickle.dump(object, file)
  137. file.flush()
  138. return file
  139. def loadIO(buffer, **kwds):
  140. """load an object that was stored with dill.temp.dumpIO
  141. buffer: buffer object
  142. >>> dumpfile = dill.temp.dumpIO([1, 2, 3, 4, 5])
  143. >>> dill.temp.loadIO(dumpfile)
  144. [1, 2, 3, 4, 5]
  145. """
  146. import dill as pickle
  147. if PY3:
  148. from io import BytesIO as StringIO
  149. else:
  150. from StringIO import StringIO
  151. value = getattr(buffer, 'getvalue', buffer) # value or buffer.getvalue
  152. if value != buffer: value = value() # buffer.getvalue()
  153. return pickle.load(StringIO(value))
  154. def dumpIO(object, **kwds):
  155. """dill.dump of object to a buffer.
  156. Loads with "dill.temp.loadIO". Returns the buffer object.
  157. >>> dumpfile = dill.temp.dumpIO([1, 2, 3, 4, 5])
  158. >>> dill.temp.loadIO(dumpfile)
  159. [1, 2, 3, 4, 5]
  160. """
  161. import dill as pickle
  162. if PY3:
  163. from io import BytesIO as StringIO
  164. else:
  165. from StringIO import StringIO
  166. file = StringIO()
  167. pickle.dump(object, file)
  168. file.flush()
  169. return file
  170. def loadIO_source(buffer, **kwds):
  171. """load an object that was stored with dill.temp.dumpIO_source
  172. buffer: buffer object
  173. alias: string name of stored object
  174. >>> f = lambda x:x**2
  175. >>> pyfile = dill.temp.dumpIO_source(f, alias='_f')
  176. >>> _f = dill.temp.loadIO_source(pyfile)
  177. >>> _f(4)
  178. 16
  179. """
  180. alias = kwds.pop('alias', None)
  181. source = getattr(buffer, 'getvalue', buffer) # source or buffer.getvalue
  182. if source != buffer: source = source() # buffer.getvalue()
  183. if PY3: source = source.decode() # buffer to string
  184. if not alias:
  185. tag = source.strip().splitlines()[-1].split()
  186. if tag[0] != '#NAME:':
  187. stub = source.splitlines()[0]
  188. raise IOError("unknown name for code: %s" % stub)
  189. alias = tag[-1]
  190. local = {}
  191. exec(source, local)
  192. _ = eval("%s" % alias, local)
  193. return _
  194. def dumpIO_source(object, **kwds):
  195. """write object source to a buffer (instead of dill.dump)
  196. Loads by with dill.temp.loadIO_source. Returns the buffer object.
  197. >>> f = lambda x:x**2
  198. >>> pyfile = dill.temp.dumpIO_source(f, alias='_f')
  199. >>> _f = dill.temp.loadIO_source(pyfile)
  200. >>> _f(4)
  201. 16
  202. Optional kwds:
  203. If 'alias' is specified, the object will be renamed to the given string.
  204. """
  205. from .source import importable, getname
  206. if PY3:
  207. from io import BytesIO as StringIO
  208. else:
  209. from StringIO import StringIO
  210. alias = kwds.pop('alias', '') #XXX: include an alias so a name is known
  211. name = str(alias) or getname(object)
  212. name = "\n#NAME: %s\n" % name
  213. #XXX: assumes kwds['dir'] is writable and on $PYTHONPATH
  214. file = StringIO()
  215. file.write(b(''.join([importable(object, alias=alias),name])))
  216. file.flush()
  217. return file
  218. del contextlib
  219. # EOF