Source code for openmdao.util.stream

import struct
import sys
import logging

try:
    import numpy
except ImportError as err:
    logging.warn("In %s: %r" % (__file__, err))

_SZ_INT = 4
_SZ_LONG = 8
_SZ_FLOAT = 4
_SZ_DOUBLE = 8

from openmdao.util.decorators import stub_if_missing_deps

@stub_if_missing_deps('numpy')
[docs]class Stream(object): """ Wrapper of standard Python :class:`file` object. Supports reading/writing int and float arrays in various formats. file_obj: file File object opened for reading or writing. binary: bool If True, the data is in binary, not text, form. big_endian: bool If True, the data bytes are in 'big-endian' order. Only meaningful if `binary`. single_precision: bool If True, floating-point data is 32 bits, not 64. Only meaningful if `binary`. integer_8: bool If True, integer data is 64 bits, not 32. Only meaningful if `binary`. unformatted: bool If True, the data is surrounded by Fortran record length markers. Only meaningful if `binary`. recordmark_8: bool If True, the record length markers are 64 bits, not 32. Only meaningful if `unformatted`. """ def __init__(self, file_obj, binary=False, big_endian=False, single_precision=False, integer_8=False, unformatted=False, recordmark_8=False): self.file = file_obj self.binary = binary if binary: self.big_endian = big_endian self.single_precision = single_precision self.integer_8 = integer_8 self.unformatted = unformatted self.recordmark_8 = recordmark_8 self.need_byteswap = False if self.big_endian: if sys.byteorder == 'little': self.need_byteswap = True else: if sys.byteorder == 'big': self.need_byteswap = True else: # Ensure sanity. self.big_endian = False self.single_precision = False self.integer_8 = False self.unformatted = False self.recordmark_8 = False self.need_byteswap = False
[docs] def close(self): """ Close underlying file. """ return self.file.close()
[docs] def reclen_ints(self, count): """ Returns record length for `count` ints. count: int Number of ints in record. """ if self.integer_8: return _SZ_LONG * count else: return _SZ_INT * count
[docs] def reclen_floats(self, count): """ Returns record length for `count` floats. count: int Number of floats in record. """ if self.single_precision: return _SZ_FLOAT * count else: return _SZ_DOUBLE * count ######## Input Operations ########
[docs] def read_int(self, full_record=False): """ Returns next integer. full_record: bool If True, then read surrounding recordmarks. Only meaningful if `unformatted`. """ if full_record and self.unformatted: reclen = self.read_recordmark() if reclen != self.reclen_ints(1): raise RuntimeError('unexpected recordlength %d' % reclen) data = self.read_ints(1)[0] if full_record and self.unformatted: reclen2 = self.read_recordmark() if reclen2 != reclen: raise RuntimeError('mismatched recordlength %d vs. %d' % (reclen2, reclen)) return int(data)
[docs] def read_ints(self, shape, order='C', full_record=False): """ Returns integers as a :mod:`numpy` array of `shape`. shape: tuple(int) Dimensions of returned array. order: string If 'C', the data is in row-major order. If 'Fortran', the data is in column-major order. full_record: bool If True, then read surrounding recordmarks. Only meaningful if `unformatted`. """ reshape = False count = 1 try: for size in shape: count *= size reshape = True except TypeError: count = shape if full_record and self.unformatted: reclen = self.read_recordmark() if reclen != self.reclen_ints(count): raise RuntimeError('unexpected recordlength %d' % reclen) sep = '' if self.binary else ' ' dtype = numpy.int64 if self.integer_8 else numpy.int32 data = numpy.fromfile(self.file, dtype=dtype, count=count, sep=sep) if self.need_byteswap: data.byteswap(True) if full_record and self.unformatted: reclen2 = self.read_recordmark() if reclen2 != reclen: raise RuntimeError('mismatched recordlength %d vs. %d' % (reclen2, reclen)) return data.reshape(shape, order=order) if reshape else data
[docs] def read_float(self, full_record=False): """ Returns next float. full_record: bool If True, then read surrounding recordmarks. Only meaningful if `unformatted`. """ if full_record and self.unformatted: reclen = self.read_recordmark() if reclen != self.reclen_floats(1): raise RuntimeError('unexpected recordlength %d' % reclen) data = self.read_floats(1)[0] if full_record and self.unformatted: reclen2 = self.read_recordmark() if reclen2 != reclen: raise RuntimeError('mismatched recordlength %d vs. %d' % (reclen2, reclen)) return float(data)
[docs] def read_floats(self, shape, order='C', full_record=False): """ Returns floats as a :mod:`numpy` array of `shape`. shape: tuple(int) Dimensions of returned array. order: string If 'C', the data is in row-major order. If 'Fortran', the data is in column-major order. full_record: bool If True, then read surrounding recordmarks. Only meaningful if `unformatted`. """ reshape = False count = 1 try: for size in shape: count *= size reshape = True except TypeError: count = shape if full_record and self.unformatted: reclen = self.read_recordmark() if reclen != self.reclen_floats(count): raise RuntimeError('unexpected recordlength %d' % reclen) sep = '' if self.binary else ' ' dtype = numpy.float32 if self.single_precision else numpy.float64 data = numpy.fromfile(self.file, dtype=dtype, count=count, sep=sep) if self.need_byteswap: data.byteswap(True) if full_record and self.unformatted: reclen2 = self.read_recordmark() if reclen2 != reclen: raise RuntimeError('mismatched recordlength %d vs. %d' % (reclen2, reclen)) return data.reshape(shape, order=order) if reshape else data
[docs] def read_recordmark(self): """ Returns value of next recordmark. """ fmt = '>' if self.big_endian else '<' fmt += 'q' if self.recordmark_8 else 'i' size = _SZ_LONG if self.recordmark_8 else _SZ_INT return struct.unpack(fmt, self.file.read(size))[0] ######## Output Operations ########
[docs] def write_int(self, value, fmt='%d', sep='', full_record=False): """ Writes an integer. value: int Value to be written. fmt: string Format to use when writing as text. sep: string Appended to stream after writing `value`. full_record: bool If True, then write surrounding recordmarks. Only meaningful if `unformatted`. """ if self.binary: if full_record and self.unformatted: self.write_recordmark(self.reclen_ints(1)) fmt = '>' if self.big_endian else '<' fmt += 'q' if self.integer_8 else 'i' self.file.write(struct.pack(fmt, value)) if full_record and self.unformatted: self.write_recordmark(self.reclen_ints(1)) else: self.file.write(fmt % value) if full_record: self.file.write('\n') elif sep: self.file.write(sep)
[docs] def write_ints(self, data, order='C', fmt='%d', sep=' ', linecount=0, full_record=False): """ Writes an integer array. data: :class:`numpy.ndarray` Integer data array. order: string If 'C', the data is written in row-major order. If 'Fortran', the data is written in column-major order. fmt: string Format specifier for each item. sep: string Separator between items. linecount: int If > zero, then at most `linecount` values are written per line. full_record: bool If True, then write surrounding recordmarks. Only meaningful if `unformatted`. """ if not isinstance(data, numpy.ndarray): data = numpy.array(data) if self.binary: if full_record and self.unformatted: self.write_recordmark(self.reclen_ints(data.size)) arr = data if self.integer_8: if data.itemsize != _SZ_LONG: arr = numpy.array(data, dtype=numpy.int64) elif data.itemsize != _SZ_INT: arr = numpy.array(data, dtype=numpy.int32) if self.need_byteswap: arr.byteswap(True) try: self.file.write(arr.tostring(order=order)) finally: if self.need_byteswap: arr.byteswap(True) if full_record and self.unformatted: self.write_recordmark(self.reclen_ints(data.size)) else: self.write_array(data, order, fmt, sep, linecount)
[docs] def write_float(self, value, fmt='%.16g', sep='', full_record=False): """ Writes a float. value: float Value to be written. fmt: string Format to use when writing as text. sep: string Appended to stream after writing `value`. full_record: bool If True, then write surrounding recordmarks. Only meaningful if `unformatted`. """ if self.binary: if full_record and self.unformatted: self.write_recordmark(self.reclen_floats(1)) fmt = '>' if self.big_endian else '<' fmt += 'f' if self.single_precision else 'd' self.file.write(struct.pack(fmt, value)) if full_record and self.unformatted: self.write_recordmark(self.reclen_floats(1)) else: self.file.write(fmt % value) if full_record: self.file.write('\n') elif sep: self.file.write(sep)
[docs] def write_floats(self, data, order='C', fmt='%.16g', sep=' ', linecount=0, full_record=False): """ Writes a float array. data: :class:`numpy.ndarray` Float data array. order: string If 'C', the data is written in row-major order. If 'Fortran', the data is written in column-major order. fmt: string Format specifier for each item. sep: string Separator between items. linecount: int If > zero, then at most `linecount` values are written per line. full_record: bool If True, then write surrounding recordmarks. Only meaningful if `unformatted`. """ if not isinstance(data, numpy.ndarray): data = numpy.array(data) if self.binary: if full_record and self.unformatted: self.write_recordmark(self.reclen_floats(data.size)) arr = data if self.single_precision: if data.itemsize != _SZ_FLOAT: arr = numpy.array(data, dtype=numpy.float32) elif data.itemsize != _SZ_DOUBLE: arr = numpy.array(data, dtype=numpy.float64) if self.need_byteswap: arr.byteswap(True) try: self.file.write(arr.tostring(order=order)) finally: if self.need_byteswap: arr.byteswap(True) if full_record and self.unformatted: self.write_recordmark(self.reclen_floats(data.size)) else: self.write_array(data, order, fmt, sep, linecount)
[docs] def write_array(self, data, order='C', fmt='%s', sep=' ', linecount=0): """ Writes array as text. data: :class:`numpy.ndarray` Data array. order: string If 'C', the data is written in row-major order. If 'Fortran', the data is written in column-major order. fmt: string Format specifier for each item. sep: string Separator between items. linecount: int If > zero, then at most `linecount` values are written per line. """ shape = data.shape indices = [0 for i in shape] _item = data.item _write = self.file.write if order == 'C': # Row-major order. sequence = range(len(shape)) sequence.reverse() elif order == 'Fortran': # Column-major order. sequence = range(len(shape)) else: raise ValueError("order must be 'C' or 'Fortran'") count = 0 while True: _write(fmt % _item(*indices)) count += 1 if linecount > 0: if count >= linecount: _write('\n') count = 0 if count: _write(sep) for i in sequence: indices[i] += 1 if indices[i] >= shape[i]: indices[i] = 0 else: break else: break if count: _write('\n')
[docs] def write_recordmark(self, length): """ Writes recordmark. length: int Length of record (bytes). """ fmt = '>' if self.big_endian else '<' fmt += 'q' if self.recordmark_8 else 'i' self.file.write(struct.pack(fmt, length))
OpenMDAO Home