# Licensed under a 3-clause BSD style license - see PYFITS.rst
from __future__ import division
from __future__ import with_statement
import gzip
import os
import sys
import tempfile
import urllib
import warnings
import zipfile
import numpy as np
from numpy import memmap as Memmap
from .util import (isreadable, iswritable, isfile, fileobj_open, fileobj_name,
fileobj_closed, fileobj_mode, _array_from_file,
_array_to_file, _write_string)
from ...utils import deprecated
# File object open modes
PYTHON_MODES = {'readonly': 'rb', 'copyonwrite': 'rb', 'update': 'rb+',
'append': 'ab+', 'ostream': 'wb', 'denywrite': 'rb'}
# readonly actually uses copyonwrite for mmap so that readonly without mmap and
# with mmap still have to same behavior with regard to updating the array. To
# get a truly readonly mmap use denywrite
# the name 'denywrite' comes from a deprecated flag to mmap() on Linux--it
# should be clarified that 'denywrite' mode is not directly analogous to the
# use of that flag; it was just taken, for lack of anything better, as a name
# that means something like "read only" but isn't readonly.
MEMMAP_MODES = {'readonly': 'c', 'copyonwrite': 'c', 'update': 'r+',
'append': 'c', 'denywrite': 'r'}
# TODO: Eventually raise a warning, and maybe even later disable the use of
# 'copyonwrite' and 'denywrite' modes unless memmap=True. For now, however,
# that would generate too many warnings for too many users. If nothing else,
# wait until the new logging system is in place.
GZIP_MAGIC = u'\x1f\x8b\x08'.encode('raw-unicode-escape')
PKZIP_MAGIC = u'\x50\x4b\x03\x04'.encode('raw-unicode-escape')
class _File(object):
"""
Represents a FITS file on disk (or in some other file-like object).
"""
def __init__(self, fileobj=None, mode='readonly', memmap=False):
if fileobj is None:
self.__file = None
self.closed = False
self.mode = mode
self.memmap = memmap
self.compression = None
self.readonly = False
self.writeonly = False
self.simulateonly = True
return
else:
self.simulateonly = False
if mode not in PYTHON_MODES:
raise ValueError("Mode '%s' not recognized" % mode)
if (isinstance(fileobj, basestring) and mode != 'append' and
not os.path.exists(fileobj) and
not os.path.splitdrive(fileobj)[0]):
#
# Not writing file and file does not exist on local machine and
# name does not begin with a drive letter (Windows), try to
# get it over the web.
#
self.name, _ = urllib.urlretrieve(fileobj)
else:
self.name = fileobj_name(fileobj)
self.closed = False
self.mode = mode
self.memmap = memmap
# Underlying fileobj is a file-like object, but an actual file object
self.file_like = False
# More defaults to be adjusted below as necessary
self.compression = None
self.readonly = False
self.writeonly = False
# Initialize the internal self.__file object
if isfile(fileobj) or isinstance(fileobj, gzip.GzipFile):
closed = fileobj_closed(fileobj)
fmode = fileobj_mode(fileobj) or PYTHON_MODES[mode]
if not closed:
# In some cases (like on Python 3) a file opened for
# appending still shows a mode of 'r+', hence the extra
# check for the append case
if ((mode == 'append' and fmode not in ('ab+', 'rb+')) or
(mode != 'append' and PYTHON_MODES[mode] != fmode)):
raise ValueError(
"Input mode '%s' (%s) does not match mode of the "
"input file (%s)." % (mode, PYTHON_MODES[mode], fmode))
self.__file = fileobj
elif isfile(fileobj):
self.__file = fileobj_open(self.name, PYTHON_MODES[mode])
# Return to the beginning of the file--in Python 3 when
# opening in append mode the file pointer is at the end of
# the file
self.__file.seek(0)
else:
self.__file = gzip.open(self.name, PYTHON_MODES[mode])
elif isinstance(fileobj, basestring):
if os.path.exists(self.name):
with fileobj_open(self.name, 'rb') as f:
magic = f.read(4)
else:
magic = ''.encode('raw-unicode-escape')
ext = os.path.splitext(self.name)[1]
if ext == '.gz' or magic.startswith(GZIP_MAGIC):
# Handle gzip files
if mode in ['update', 'append']:
raise IOError(
"Writing to gzipped fits files is not currently "
"supported")
self.__file = gzip.open(self.name)
self.compression = 'gzip'
elif ext == '.zip' or magic.startswith(PKZIP_MAGIC):
# Handle zip files
if mode in ['update', 'append']:
raise IOError(
"Writing to zipped fits files is not currently "
"supported")
zfile = zipfile.ZipFile(self.name)
namelist = zfile.namelist()
if len(namelist) != 1:
raise IOError(
"Zip files with multiple members are not supported.")
self.__file = tempfile.NamedTemporaryFile(suffix='.fits')
self.__file.write(zfile.read(namelist[0]))
zfile.close()
self.compression = 'zip'
else:
self.__file = fileobj_open(self.name, PYTHON_MODES[mode])
# Make certain we're back at the beginning of the file
self.__file.seek(0)
else:
# We are dealing with a file like object.
# Assume it is open.
self.file_like = True
self.__file = fileobj
# If there is not seek or tell methods then set the mode to
# output streaming.
if (not hasattr(self.__file, 'seek') or
not hasattr(self.__file, 'tell')):
self.mode = mode = 'ostream'
if (self.mode in ('copyonwrite', 'update', 'append') and
not hasattr(self.__file, 'write')):
raise IOError("File-like object does not have a 'write' "
"method, required for mode '%s'."
% self.mode)
if (self.mode in ('readonly', 'denywrite') and
not hasattr(self.__file, 'read')):
raise IOError("File-like object does not have a 'read' "
"method, required for mode %r."
% self.mode)
if isinstance(fileobj, gzip.GzipFile):
self.compression = 'gzip'
elif isinstance(fileobj, zipfile.ZipFile):
# Reading from zip files is supported but not writing (yet)
self.compression = 'zip'
if (mode in ('readonly', 'copyonwrite', 'denywrite') or
(self.compression and mode == 'update')):
self.readonly = True
elif (mode == 'ostream' or
(self.compression and mode == 'append')):
self.writeonly = True
# For 'ab+' mode, the pointer is at the end after the open in
# Linux, but is at the beginning in Solaris.
if (mode == 'ostream' or self.compression or
not hasattr(self.__file, 'seek')):
# For output stream start with a truncated file.
# For compressed files we can't really guess at the size
self.size = 0
else:
pos = self.__file.tell()
self.__file.seek(0, 2)
self.size = self.__file.tell()
self.__file.seek(pos)
if self.memmap and not isfile(self.__file):
self.memmap = False
def __repr__(self):
return '<%s.%s %s>' % (self.__module__, self.__class__.__name__,
self.__file)
# Support the 'with' statement
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
@deprecated('3.0', message='This method should not be treated as public.')
def getfile(self):
"""Will be going away as soon as I figure out how."""
return self.__file
def readable(self):
if self.writeonly:
return False
return isreadable(self.__file)
def read(self, size=None):
if not hasattr(self.__file, 'read'):
raise EOFError
return self.__file.read(size)
def readarray(self, size=None, offset=0, dtype=np.uint8, shape=None):
"""
Similar to file.read(), but returns the contents of the underlying
file as a numpy array (or mmap'd array if memmap=True) rather than a
string.
Usually it's best not to use the `size` argument with this method, but
it's provided for compatibility.
"""
if not hasattr(self.__file, 'read'):
raise EOFError
if not isinstance(dtype, np.dtype):
dtype = np.dtype(dtype)
if size and size % dtype.itemsize != 0:
raise ValueError('size %d not a multiple of %s' % (size, dtype))
if isinstance(shape, int):
shape = (shape,)
if size and shape:
actualsize = sum(dim * dtype.itemsize for dim in shape)
if actualsize < size:
raise ValueError('size %d is too few bytes for a %s array of '
'%s' % (size, shape, dtype))
if actualsize < size:
raise ValueError('size %d is too many bytes for a %s array of '
'%s' % (size, shape, dtype))
if size and not shape:
shape = (size // dtype.itemsize,)
if not (size or shape):
warnings.warn('No size or shape given to readarray(); assuming a '
'shape of (1,)')
shape = (1,)
if self.memmap:
return Memmap(self.__file, offset=offset,
mode=MEMMAP_MODES[self.mode], dtype=dtype,
shape=shape).view(np.ndarray)
else:
count = reduce(lambda x, y: x * y, shape)
pos = self.__file.tell()
self.__file.seek(offset)
data = _array_from_file(self.__file, dtype, count, '')
data.shape = shape
self.__file.seek(pos)
return data
def writable(self):
if self.readonly:
return False
return iswritable(self.__file)
def write(self, string):
if hasattr(self.__file, 'write'):
_write_string(self.__file, string)
def writearray(self, array):
"""
Similar to file.write(), but writes a numpy array instead of a string.
Also like file.write(), a flush() or close() may be needed before
the file on disk reflects the data written.
"""
if hasattr(self.__file, 'write'):
_array_to_file(array, self.__file)
def flush(self):
if hasattr(self.__file, 'flush'):
self.__file.flush()
def seek(self, offset, whence=0):
# In newer Python versions, GzipFiles support the whence argument, but
# I don't think it was added until 2.6; instead of assuming it's
# present, we implement our own support for it here
if not hasattr(self.__file, 'seek'):
return
if isinstance(self.__file, gzip.GzipFile):
if whence:
if whence == 1:
offset = self.__file.offset + offset
else:
raise ValueError('Seek from end not supported')
self.__file.seek(offset)
else:
self.__file.seek(offset, whence)
pos = self.__file.tell()
if self.size and pos > self.size:
warnings.warn('File may have been truncated: actual file length '
'(%i) is smaller than the expected size (%i)' %
(self.size, pos))
def tell(self):
if not hasattr(self.__file, 'tell'):
raise EOFError
return self.__file.tell()
def truncate(self, size=None):
if hasattr(self.__file, 'truncate'):
self.__file.truncate(size)
def close(self):
"""
Close the 'physical' FITS file.
"""
if hasattr(self.__file, 'close'):
self.__file.close()
self.closed = True