PageRenderTime 77ms CodeModel.GetById 41ms RepoModel.GetById 1ms app.codeStats 0ms

/astropy/io/fits/file.py

https://github.com/jiffyclub/astropy
Python | 344 lines | 264 code | 35 blank | 45 comment | 41 complexity | fb5df0f7f42d2fd56d522f352883a48f MD5 | raw file
  1. # Licensed under a 3-clause BSD style license - see PYFITS.rst
  2. from __future__ import division
  3. from __future__ import with_statement
  4. import gzip
  5. import os
  6. import sys
  7. import tempfile
  8. import urllib
  9. import warnings
  10. import zipfile
  11. import numpy as np
  12. from numpy import memmap as Memmap
  13. from .util import (isreadable, iswritable, isfile, fileobj_open, fileobj_name,
  14. fileobj_closed, fileobj_mode, _array_from_file,
  15. _array_to_file, _write_string)
  16. from ...utils import deprecated
  17. # File object open modes
  18. PYTHON_MODES = {'readonly': 'rb', 'copyonwrite': 'rb', 'update': 'rb+',
  19. 'append': 'ab+', 'ostream': 'wb', 'denywrite': 'rb'}
  20. # readonly actually uses copyonwrite for mmap so that readonly without mmap and
  21. # with mmap still have to same behavior with regard to updating the array. To
  22. # get a truly readonly mmap use denywrite
  23. # the name 'denywrite' comes from a deprecated flag to mmap() on Linux--it
  24. # should be clarified that 'denywrite' mode is not directly analogous to the
  25. # use of that flag; it was just taken, for lack of anything better, as a name
  26. # that means something like "read only" but isn't readonly.
  27. MEMMAP_MODES = {'readonly': 'c', 'copyonwrite': 'c', 'update': 'r+',
  28. 'append': 'c', 'denywrite': 'r'}
  29. # TODO: Eventually raise a warning, and maybe even later disable the use of
  30. # 'copyonwrite' and 'denywrite' modes unless memmap=True. For now, however,
  31. # that would generate too many warnings for too many users. If nothing else,
  32. # wait until the new logging system is in place.
  33. GZIP_MAGIC = u'\x1f\x8b\x08'.encode('raw-unicode-escape')
  34. PKZIP_MAGIC = u'\x50\x4b\x03\x04'.encode('raw-unicode-escape')
  35. class _File(object):
  36. """
  37. Represents a FITS file on disk (or in some other file-like object).
  38. """
  39. def __init__(self, fileobj=None, mode='readonly', memmap=False):
  40. if fileobj is None:
  41. self.__file = None
  42. self.closed = False
  43. self.mode = mode
  44. self.memmap = memmap
  45. self.compression = None
  46. self.readonly = False
  47. self.writeonly = False
  48. self.simulateonly = True
  49. return
  50. else:
  51. self.simulateonly = False
  52. if mode not in PYTHON_MODES:
  53. raise ValueError("Mode '%s' not recognized" % mode)
  54. if (isinstance(fileobj, basestring) and mode != 'append' and
  55. not os.path.exists(fileobj) and
  56. not os.path.splitdrive(fileobj)[0]):
  57. #
  58. # Not writing file and file does not exist on local machine and
  59. # name does not begin with a drive letter (Windows), try to
  60. # get it over the web.
  61. #
  62. self.name, _ = urllib.urlretrieve(fileobj)
  63. else:
  64. self.name = fileobj_name(fileobj)
  65. self.closed = False
  66. self.mode = mode
  67. self.memmap = memmap
  68. # Underlying fileobj is a file-like object, but an actual file object
  69. self.file_like = False
  70. # More defaults to be adjusted below as necessary
  71. self.compression = None
  72. self.readonly = False
  73. self.writeonly = False
  74. # Initialize the internal self.__file object
  75. if isfile(fileobj) or isinstance(fileobj, gzip.GzipFile):
  76. closed = fileobj_closed(fileobj)
  77. fmode = fileobj_mode(fileobj) or PYTHON_MODES[mode]
  78. if not closed:
  79. # In some cases (like on Python 3) a file opened for
  80. # appending still shows a mode of 'r+', hence the extra
  81. # check for the append case
  82. if ((mode == 'append' and fmode not in ('ab+', 'rb+')) or
  83. (mode != 'append' and PYTHON_MODES[mode] != fmode)):
  84. raise ValueError(
  85. "Input mode '%s' (%s) does not match mode of the "
  86. "input file (%s)." % (mode, PYTHON_MODES[mode], fmode))
  87. self.__file = fileobj
  88. elif isfile(fileobj):
  89. self.__file = fileobj_open(self.name, PYTHON_MODES[mode])
  90. # Return to the beginning of the file--in Python 3 when
  91. # opening in append mode the file pointer is at the end of
  92. # the file
  93. self.__file.seek(0)
  94. else:
  95. self.__file = gzip.open(self.name, PYTHON_MODES[mode])
  96. elif isinstance(fileobj, basestring):
  97. if os.path.exists(self.name):
  98. with fileobj_open(self.name, 'rb') as f:
  99. magic = f.read(4)
  100. else:
  101. magic = ''.encode('raw-unicode-escape')
  102. ext = os.path.splitext(self.name)[1]
  103. if ext == '.gz' or magic.startswith(GZIP_MAGIC):
  104. # Handle gzip files
  105. if mode in ['update', 'append']:
  106. raise IOError(
  107. "Writing to gzipped fits files is not currently "
  108. "supported")
  109. self.__file = gzip.open(self.name)
  110. self.compression = 'gzip'
  111. elif ext == '.zip' or magic.startswith(PKZIP_MAGIC):
  112. # Handle zip files
  113. if mode in ['update', 'append']:
  114. raise IOError(
  115. "Writing to zipped fits files is not currently "
  116. "supported")
  117. zfile = zipfile.ZipFile(self.name)
  118. namelist = zfile.namelist()
  119. if len(namelist) != 1:
  120. raise IOError(
  121. "Zip files with multiple members are not supported.")
  122. self.__file = tempfile.NamedTemporaryFile(suffix='.fits')
  123. self.__file.write(zfile.read(namelist[0]))
  124. zfile.close()
  125. self.compression = 'zip'
  126. else:
  127. self.__file = fileobj_open(self.name, PYTHON_MODES[mode])
  128. # Make certain we're back at the beginning of the file
  129. self.__file.seek(0)
  130. else:
  131. # We are dealing with a file like object.
  132. # Assume it is open.
  133. self.file_like = True
  134. self.__file = fileobj
  135. # If there is not seek or tell methods then set the mode to
  136. # output streaming.
  137. if (not hasattr(self.__file, 'seek') or
  138. not hasattr(self.__file, 'tell')):
  139. self.mode = mode = 'ostream'
  140. if (self.mode in ('copyonwrite', 'update', 'append') and
  141. not hasattr(self.__file, 'write')):
  142. raise IOError("File-like object does not have a 'write' "
  143. "method, required for mode '%s'."
  144. % self.mode)
  145. if (self.mode in ('readonly', 'denywrite') and
  146. not hasattr(self.__file, 'read')):
  147. raise IOError("File-like object does not have a 'read' "
  148. "method, required for mode %r."
  149. % self.mode)
  150. if isinstance(fileobj, gzip.GzipFile):
  151. self.compression = 'gzip'
  152. elif isinstance(fileobj, zipfile.ZipFile):
  153. # Reading from zip files is supported but not writing (yet)
  154. self.compression = 'zip'
  155. if (mode in ('readonly', 'copyonwrite', 'denywrite') or
  156. (self.compression and mode == 'update')):
  157. self.readonly = True
  158. elif (mode == 'ostream' or
  159. (self.compression and mode == 'append')):
  160. self.writeonly = True
  161. # For 'ab+' mode, the pointer is at the end after the open in
  162. # Linux, but is at the beginning in Solaris.
  163. if (mode == 'ostream' or self.compression or
  164. not hasattr(self.__file, 'seek')):
  165. # For output stream start with a truncated file.
  166. # For compressed files we can't really guess at the size
  167. self.size = 0
  168. else:
  169. pos = self.__file.tell()
  170. self.__file.seek(0, 2)
  171. self.size = self.__file.tell()
  172. self.__file.seek(pos)
  173. if self.memmap and not isfile(self.__file):
  174. self.memmap = False
  175. def __repr__(self):
  176. return '<%s.%s %s>' % (self.__module__, self.__class__.__name__,
  177. self.__file)
  178. # Support the 'with' statement
  179. def __enter__(self):
  180. return self
  181. def __exit__(self, type, value, traceback):
  182. self.close()
  183. @deprecated('3.0', message='This method should not be treated as public.')
  184. def getfile(self):
  185. """Will be going away as soon as I figure out how."""
  186. return self.__file
  187. def readable(self):
  188. if self.writeonly:
  189. return False
  190. return isreadable(self.__file)
  191. def read(self, size=None):
  192. if not hasattr(self.__file, 'read'):
  193. raise EOFError
  194. return self.__file.read(size)
  195. def readarray(self, size=None, offset=0, dtype=np.uint8, shape=None):
  196. """
  197. Similar to file.read(), but returns the contents of the underlying
  198. file as a numpy array (or mmap'd array if memmap=True) rather than a
  199. string.
  200. Usually it's best not to use the `size` argument with this method, but
  201. it's provided for compatibility.
  202. """
  203. if not hasattr(self.__file, 'read'):
  204. raise EOFError
  205. if not isinstance(dtype, np.dtype):
  206. dtype = np.dtype(dtype)
  207. if size and size % dtype.itemsize != 0:
  208. raise ValueError('size %d not a multiple of %s' % (size, dtype))
  209. if isinstance(shape, int):
  210. shape = (shape,)
  211. if size and shape:
  212. actualsize = sum(dim * dtype.itemsize for dim in shape)
  213. if actualsize < size:
  214. raise ValueError('size %d is too few bytes for a %s array of '
  215. '%s' % (size, shape, dtype))
  216. if actualsize < size:
  217. raise ValueError('size %d is too many bytes for a %s array of '
  218. '%s' % (size, shape, dtype))
  219. if size and not shape:
  220. shape = (size // dtype.itemsize,)
  221. if not (size or shape):
  222. warnings.warn('No size or shape given to readarray(); assuming a '
  223. 'shape of (1,)')
  224. shape = (1,)
  225. if self.memmap:
  226. return Memmap(self.__file, offset=offset,
  227. mode=MEMMAP_MODES[self.mode], dtype=dtype,
  228. shape=shape).view(np.ndarray)
  229. else:
  230. count = reduce(lambda x, y: x * y, shape)
  231. pos = self.__file.tell()
  232. self.__file.seek(offset)
  233. data = _array_from_file(self.__file, dtype, count, '')
  234. data.shape = shape
  235. self.__file.seek(pos)
  236. return data
  237. def writable(self):
  238. if self.readonly:
  239. return False
  240. return iswritable(self.__file)
  241. def write(self, string):
  242. if hasattr(self.__file, 'write'):
  243. _write_string(self.__file, string)
  244. def writearray(self, array):
  245. """
  246. Similar to file.write(), but writes a numpy array instead of a string.
  247. Also like file.write(), a flush() or close() may be needed before
  248. the file on disk reflects the data written.
  249. """
  250. if hasattr(self.__file, 'write'):
  251. _array_to_file(array, self.__file)
  252. def flush(self):
  253. if hasattr(self.__file, 'flush'):
  254. self.__file.flush()
  255. def seek(self, offset, whence=0):
  256. # In newer Python versions, GzipFiles support the whence argument, but
  257. # I don't think it was added until 2.6; instead of assuming it's
  258. # present, we implement our own support for it here
  259. if not hasattr(self.__file, 'seek'):
  260. return
  261. if isinstance(self.__file, gzip.GzipFile):
  262. if whence:
  263. if whence == 1:
  264. offset = self.__file.offset + offset
  265. else:
  266. raise ValueError('Seek from end not supported')
  267. self.__file.seek(offset)
  268. else:
  269. self.__file.seek(offset, whence)
  270. pos = self.__file.tell()
  271. if self.size and pos > self.size:
  272. warnings.warn('File may have been truncated: actual file length '
  273. '(%i) is smaller than the expected size (%i)' %
  274. (self.size, pos))
  275. def tell(self):
  276. if not hasattr(self.__file, 'tell'):
  277. raise EOFError
  278. return self.__file.tell()
  279. def truncate(self, size=None):
  280. if hasattr(self.__file, 'truncate'):
  281. self.__file.truncate(size)
  282. def close(self):
  283. """
  284. Close the 'physical' FITS file.
  285. """
  286. if hasattr(self.__file, 'close'):
  287. self.__file.close()
  288. self.closed = True