def _prewriteto(self, checksum=False, inplace=False): if self._data_loaded and self.data is not None: self.data._scale_back() # check TFIELDS and NAXIS2 self._header['TFIELDS'] = len(self.data._coldefs) self._header['NAXIS2'] = self.data.shape[0] # calculate PCOUNT, for variable length tables tbsize = self.header['NAXIS1'] * self.header['NAXIS2'] heapstart = self.header.get('THEAP', tbsize) self.data._gap = heapstart - tbsize pcount = self.data._heapsize + self.data._gap if pcount > 0: self.header['PCOUNT'] = pcount # update TFORM for variable length columns for idx in range(self.data._nfields): format = self.data._coldefs.formats[idx] if isinstance(format, _FormatP): max = self.data.field(idx).max format = _FormatP(format, repeat=format.repeat, max=max) self._header['TFORM' + str(idx + 1)] = format.tform return super(_TableBaseHDU, self)._prewriteto(checksum, inplace)
def _load_data(cls, fileobj, coldefs=None): """ Read the table data from the ASCII file output by BinTableHDU.dump(). """ close_file = False if isinstance(fileobj, basestring): fileobj = open(fileobj, 'r') close_file = True initialpos = fileobj.tell() # We'll be returning here later linereader = csv.reader(fileobj, dialect=FITSTableDumpDialect) # First we need to do some preprocessing on the file to find out how # much memory we'll need to reserve for the table. This is necessary # even if we already have the coldefs in order to determine how many # rows to reserve memory for vla_lengths = [] recformats = [] names = [] nrows = 0 if coldefs is not None: recformats = coldefs._recformats names = coldefs.names def update_recformats(value, idx): fitsformat = _scalar_to_format(value) recformat = _convert_format(fitsformat) if idx >= len(recformats): recformats.append(recformat) else: if _cmp_recformats(recformats[idx], recformat) < 0: recformats[idx] = recformat # TODO: The handling of VLAs could probably be simplified a bit for row in linereader: nrows += 1 if coldefs is not None: continue col = 0 idx = 0 while idx < len(row): if row[idx] == 'VLA_Length=': if col < len(vla_lengths): vla_length = vla_lengths[col] else: vla_length = int(row[idx + 1]) vla_lengths.append(vla_length) idx += 2 while vla_length: update_recformats(row[idx], col) vla_length -= 1 idx += 1 col += 1 else: if col >= len(vla_lengths): vla_lengths.append(None) update_recformats(row[idx], col) col += 1 idx += 1 # Update the recformats for any VLAs for idx, length in enumerate(vla_lengths): if length is not None: recformats[idx] = str(length) + recformats[idx] dtype = np.rec.format_parser(recformats, names, None).dtype # TODO: In the future maybe enable loading a bit at a time so that we # can convert from this format to an actual FITS file on disk without # needing enough physical memory to hold the entire thing at once; # new_table() could use a similar feature. hdu = new_table(np.recarray(shape=1, dtype=dtype), nrows=nrows, fill=True) data = hdu.data for idx, length in enumerate(vla_lengths): if length is not None: arr = data.columns._arrays[idx] dt = recformats[idx][len(str(length)):] recformats[idx] = _FormatP(dt, max=length) data.columns._recformats[idx] = recformats[idx] data._convert[idx] = _makep(arr, arr, recformats[idx]) # Jump back to the start of the data and create a new line reader fileobj.seek(initialpos) linereader = csv.reader(fileobj, dialect=FITSTableDumpDialect) for row, line in enumerate(linereader): col = 0 idx = 0 while idx < len(line): if line[idx] == 'VLA_Length=': vla_len = vla_lengths[col] idx += 2 data[row][col][:] = line[idx:idx + vla_len] idx += vla_len else: # TODO: This won't work for complex-valued types; fix this # Kind of silly special handling for bools val = line[idx] if recformats[col] == FITS2NUMPY['L']: val = bool(int(val)) elif recformats[col] == FITS2NUMPY['M']: # For some reason, in arrays/fields where numpy expects # a complex it's not happy to take a string # representation (though it's happy to do that in other # contexts), so we have to convert the string # representation for it: val = complex(val) data[row][col] = val idx += 1 col += 1 if close_file: fileobj.close() return data