import os import numpy import numpy.ctypeslib import ctypes as CT import umfile _len_real_hdr = 19 _len_int_hdr = 45 class File_type(CT.Structure): _fields_ = [("format", CT.c_int), ("byte_ordering", CT.c_int), ("word_size", CT.c_int)] def _get_ctypes_array(dtype, size=None): """ get ctypes corresponding to a numpy array of a given type; the size should not be necessary unless the storage for the array is allocated in the C code """ kwargs = {'dtype': dtype, 'ndim': 1, 'flags': ('C_CONTIGUOUS', 'WRITEABLE')} if size: kwargs['shape'] = (size,) return numpy.ctypeslib.ndpointer(**kwargs) def _gen_rec_class(int_type, float_type): class Rec(CT.Structure): """ ctypes object corresponding to the Rec object in the C code, """ _fields_ = [("int_hdr", _get_ctypes_array(int_type, _len_int_hdr)), ("real_hdr", _get_ctypes_array(float_type, _len_real_hdr)), ("header_offset", CT.c_size_t), ("data_offset", CT.c_size_t), ("disk_length", CT.c_size_t), ("_internp", CT.c_void_p)] return Rec Rec32 = _gen_rec_class(numpy.int32, numpy.float32) Rec64 = _gen_rec_class(numpy.int64, numpy.float64) def _gen_var_class(rec_class): class Var(CT.Structure): """ ctypes object corresponding to the Var object in the C code """ _fields_ = [("recs", CT.POINTER(CT.POINTER(rec_class))), ("nz", CT.c_int), ("nt", CT.c_int), ("supervar_index", CT.c_int), ("_internp", CT.c_void_p)] return Var Var32 = _gen_var_class(Rec32) Var64 = _gen_var_class(Rec64) def _gen_file_class(var_class): class File(CT.Structure): """ ctypes object corresponding to the File object in the C code """ _fields_ = [("fd", CT.c_int), ("file_type", File_type), ("nvars", CT.c_int), ("vars", CT.POINTER(CT.POINTER(var_class))), ("_internp", CT.c_void_p)] return File File32 = _gen_file_class(Var32) File64 = _gen_file_class(Var64) class Enum(object): def __init__(self, *names): self.names = names def as_name(self, val): if isinstance(val, str): return val else: return self.names[val] def as_index(self, val): if isinstance(val, int): return val return self.names.index(val) enum_file_format = Enum("PP", "FF") enum_byte_ordering = Enum("little_endian", "big_endian") enum_data_type = Enum("integer", "real") class CInterface(object): """ Interface to the C shared library functions """ def __init__(self, lib_name = "umfile.so"): lib_dir = os.path.join(os.path.dirname(__file__) or ".", "c-lib") lib_path = os.path.join(lib_dir, lib_name) self.lib = CT.CDLL(lib_path) def _is_null_pointer(self, ptr): try: ptr.contents return False except ValueError: return True def detect_file_type(self, fd): """ auto-detect file type; returns a File_type ctypes object that can be passed to file_parse(), or raises an exception if file type cannot be detected """ file_type = File_type() rv = self.lib.detect_file_type(fd, CT.pointer(file_type)) if rv != 0: raise umfile.UMFileException("File type could not be detected") return file_type def file_type_obj_to_dict(self, file_type): """ converts a FileType object returned by detect_file_type() into a dictionary that include meaningful string values in place of the integers that derive from the C enum statments, specifically: 'format': 'PP' or 'FF' 'byte_ordering': 'little_endian' or 'big_endian' and also 'word_size': 4 or 8 """ format = enum_file_format.as_name(file_type.format) byte_ordering = enum_byte_ordering.as_name(file_type.byte_ordering) word_size = file_type.word_size return { 'format': format, 'byte_ordering': byte_ordering, 'word_size': word_size } def create_file_type(self, format, byte_ordering, word_size): """ takes string input values: 'format': 'PP' or 'FF' 'byte_ordering': 'little_endian' or 'big_endian' 'word_size': 4 or 8 and returns a FileType object (ctypes structure containing integer values) that can be passed to file_parse() """ return File_type(format = enum_file_format.as_index(format), byte_ordering = enum_byte_ordering.as_index(byte_ordering), word_size = word_size) def set_word_size(self, val): """ Sets the word size used to interpret returned pointers from subsequent calls, in particular the pointers to PP headers embedded in the tree of objects returned by file_parse() and the data array that is populated by read_record_data(). the 'val' argument contains either just the word_size value to use or a file_type object from which it is to be extracted """ if isinstance(val, File_type): word_size = val.word_size else: word_size = val if word_size == 4: self.file_class = File32 self.file_data_int_type = numpy.int32 self.file_data_real_type = numpy.float32 elif word_size == 8: self.file_class = File64 self.file_data_int_type = numpy.int64 self.file_data_real_type = numpy.float64 else: raise ValueError("word size must be 4 or 8 (not %s)" % word_size) def _get_ctypes_int_array(self, size=None): return _get_ctypes_array(self.file_data_int_type, size) def _get_ctypes_real_array(self, size=None): return _get_ctypes_array(self.file_data_real_type, size) def _get_empty_real_array(self, size): """ get empty numpy real array according to word size previously set with set_word_size() """ return numpy.empty(size, dtype = self.file_data_real_type) def _get_empty_int_array(self, size): """ as _get_empty_real_array but for int """ return numpy.empty(size, dtype = self.file_data_int_type) def parse_file(self, fh, file_type): """ Given an open file handle, work out information from the file, and return this in a dictionary, of which currently the only key actually implemented is 'vars', containing a list of variables, as that is all that the caller requires. arguments: fh - low-level file handle (integer) file_type - File_type object as returned by detect_file_type() or create_file_type() """ func = self.lib.file_parse file_p_type = CT.POINTER(self.file_class) func.restype = file_p_type file_p = func(fh, file_type) if self._is_null_pointer(file_p): raise umfile.UMFileException("file parsing failed") file = file_p.contents c_vars = file.vars[:file.nvars] rv = {'vars': map(self.c_var_to_py_var, c_vars)} # now that we have copied all the data into python objects for the # caller, free any memory allocated in the C code before returning free_func = self.lib.file_free free_func._fields_ = file_p_type free_func(file_p) return rv def c_var_to_py_var(self, c_var_p): """ create a umfile.Var object from a ctypes object corresponding to 'Var*' in the C code """ c_var = c_var_p.contents nz = c_var.nz nt = c_var.nt svi = c_var.supervar_index if svi < 0: svi = None c_recs = c_var.recs recs = [self.c_rec_to_py_rec(c_recs[recid]) for recid in range(nz * nt)] return umfile.Var(recs, nz, nt, svi) def c_rec_to_py_rec(self, c_rec_p): """ create a umfile.Rec object from a ctypes object corresponding to 'Rec*' in the C code """ c_rec = c_rec_p.contents # numpy.copy used here so we can go back and free the memory allocated # by C without affecting the python object int_hdr = numpy.copy(numpy.ctypeslib.as_array(c_rec.int_hdr)) real_hdr = numpy.copy(numpy.ctypeslib.as_array(c_rec.real_hdr)) header_offset = c_rec.header_offset data_offset = c_rec.data_offset disk_length = c_rec.disk_length return umfile.Rec(int_hdr, real_hdr, header_offset, data_offset, disk_length) def get_type_and_num_words(self, int_hdr): """ from integer header, work out data type and number of words to read (read_record_data requires this) returns 2-tuple of: data type: 'integer' or 'real' number of words """ word_size = int_hdr.itemsize self.lib.get_type_and_num_words.argtypes = [ CT.c_int, self._get_ctypes_int_array(), CT.POINTER(CT.c_int), CT.POINTER(CT.c_size_t) ] data_type = CT.c_int() num_words = CT.c_size_t() rv = self.lib.get_type_and_num_words(word_size, int_hdr, CT.pointer(data_type), CT.pointer(num_words)) if rv != 0: raise umfile.UMFileException("error determining data type and size from integer header") return enum_data_type.as_name(data_type.value), num_words.value def get_extra_data_offset_and_length(self, int_hdr, data_offset, disk_length): """ from integer header, gets offset and length of extra data returns 2-tuple (offset, length), both in units of BYTES """ word_size = int_hdr.itemsize func = self.lib.get_extra_data_offset_and_length func.argtypes = [ CT.c_int, self._get_ctypes_int_array(), CT.c_size_t, CT.c_size_t, CT.POINTER(CT.c_size_t), CT.POINTER(CT.c_size_t) ] extra_data_offset = CT.c_size_t() extra_data_length = CT.c_size_t() rv = func(word_size, int_hdr, data_offset, disk_length, CT.pointer(extra_data_offset), CT.pointer(extra_data_length)) if rv != 0: raise umfile.UMFileException("error determining extra data length from integer header") return extra_data_offset.value, extra_data_length.value def read_header(self, fd, header_offset, byte_ordering, word_size): """ reads header from open file, returning as 2-tuple (int_hdr, real_hdr) of numpy arrays """ self.lib.read_header.argtypes = [ CT.c_int, CT.c_size_t, CT.c_int, CT.c_int, self._get_ctypes_int_array(), self._get_ctypes_real_array()] int_hdr = self._get_empty_int_array(_len_int_hdr) real_hdr = self._get_empty_real_array(_len_real_hdr) rv = self.lib.read_header(fd, header_offset, enum_byte_ordering.as_index(byte_ordering), word_size, int_hdr, real_hdr) if rv != 0: raise umfile.UMFileException("error reading header data") return int_hdr, real_hdr def read_extra_data(self, fd, extra_data_offset, extra_data_length, byte_ordering, word_size): """ reads record data from open file inputs: fd - integer low-level file descriptor extra_data_offset - offset in bytes extra_disk_length - disk length of extra data in bytes byte_ordering - 'little_endian' or 'big_endian' word_size - 4 or 8 returns: extra data as string) """ extra_data = "\0" * extra_data_length self.lib.read_extra_data.argtypes = [ CT.c_int, CT.c_size_t, CT.c_size_t, CT.c_int, CT.c_int, CT.c_char_p ] rv = self.lib.read_extra_data(fd, extra_data_offset, extra_data_length, enum_byte_ordering.as_index(byte_ordering), word_size, extra_data) if rv != 0: raise umfile.UMFileException("error reading extra data") return extra_data def read_record_data(self, fd, data_offset, disk_length, byte_ordering, word_size, int_hdr, real_hdr, data_type, nwords): """ reads record data from open file inputs: fd - integer low-level file descriptor data_offset - offset in words disk_length - disk length of data record in words byte_ordering - 'little_endian' or 'big_endian' word_size - 4 or 8 int_hdr - integer PP headers (numpy array) real_hdr - real PP headers (numpy array) data_type - 'integer' or 'real' nwords - number of words to read type and nwords should have been returned by get_type_and_num_words() """ if data_type == 'integer': data = self._get_empty_int_array(nwords) ctypes_data = self._get_ctypes_int_array() elif data_type == 'real': data = self._get_empty_real_array(nwords) ctypes_data = self._get_ctypes_real_array() else: raise ValueError("data_type must be 'integer' or 'real'") self.lib.read_record_data.argtypes = [ CT.c_int, CT.c_size_t, CT.c_size_t, CT.c_int, CT.c_int, self._get_ctypes_int_array(), self._get_ctypes_real_array(), CT.c_size_t, ctypes_data ] rv = self.lib.read_record_data(fd, data_offset, disk_length, enum_byte_ordering.as_index(byte_ordering), word_size, int_hdr, real_hdr, nwords, data) if rv != 0: raise umfile.UMFileException("error reading record data") return data if __name__ == "__main__": import sys c = CInterface() fd = os.open(sys.argv[1], os.O_RDONLY) file_type = c.detect_file_type(fd) print c.file_type_obj_to_dict(file_type) c.set_word_size(file_type) info = c.parse_file(fd, file_type) for var in info['vars']: print "nz = %s, nt = %s" % (var.nz, var.nt) for rec in var.recs: print rec.hdr_offset print "data offset", rec.data_offset print "disk length", rec.disk_length print "int hdr", rec.int_hdr print "real hdr", rec.real_hdr data_type, nwords = c.get_type_and_num_words(rec.int_hdr) print "data_type = %s nwords = %s" % (data_type, nwords) data = c.read_record_data(fd, rec.data_offset, rec.disk_length, file_type.byte_ordering, file_type.word_size, rec.int_hdr, rec.real_hdr, data_type, nwords) print "data (%s values): %s ... %s" % (len(data), data[:10], data[-10:]) extra_data_offset, extra_data_length = \ c.get_extra_data_offset_and_length(rec.int_hdr, rec.data_offset, rec.disk_length) print "extra data offset: %s" % extra_data_offset print "extra data length: %s" % extra_data_length extra_data = c.read_extra_data(fd, extra_data_offset, extra_data_length, file_type.byte_ordering, file_type.word_size) print "extra data (%s bytes) read" % (len(extra_data)) print c.read_header(fd, info['vars'][0].recs[0].hdr_offset, file_type.byte_ordering, file_type.word_size)