1 """
2 Writing data in FITS binary tables
3 """
4
5
6
7
8
9
10
11 import os
12 import tempfile
13 import time
14
15 import numpy
16
17 from gavo import base
18 from gavo import rsc
19 from gavo import utils
20 from gavo.formats import common
21 from gavo.utils import pyfits
22 from gavo.votable import common as votcommon
23
24
25 _fitsCodeMap = {
26 "short": "I",
27 "int": "J",
28 "long": "K",
29 "float": "E",
30 "double": "D",
31 "boolean": "L",
32 "char": "A",
33 "unicodeChar": "A",
34 }
35
36 _typeConstructors = {
37 "short": int,
38 "int": int,
39 "long": int,
40 "float": float,
41 "double": float,
42 "boolean": int,
43 "char": str,
44 "unicodeChar": str,
45 }
46
47
49 """returns a pyfits-capable column array for strings stored in the colInd-th
50 column of values.
51 """
52 try:
53 arr = numpy.array([str(v[colInd]) for v in values], dtype=numpy.str)
54 except UnicodeEncodeError:
55
56 def _(s):
57 if s is None:
58 return ""
59 else:
60 return s.encode("utf-8")
61
62 arr = numpy.array([_(v[colInd]) for v in values], dtype=numpy.str)
63 return "%dA"%arr.itemsize, arr
64
65
67 """returns a null value we consider ok for a column described by colDesc.
68
69 This is supposed to be in the column data type.
70 """
71 if votcommon.getLength(colDesc["arraysize"])>1:
72
73 return None
74
75 nullValue = colDesc["nullvalue"]
76 if nullValue is None:
77
78 if (colDesc["datatype"]=="float"
79 or colDesc["datatype"]=="double"):
80 nullValue = float("NaN")
81 elif colDesc["datatype"]=="text":
82 nullValue = ""
83 else:
84 nullValue = _typeConstructors[colDesc["datatype"]](nullValue)
85
86 return nullValue
87
88
90 """returns a FITS type code suitable for colDesc.
91
92 This is deficient in that we don't do variable length arrays
93 and the like for now. This may need extension when we want
94 to get fancy with FITS tables.
95 """
96 baseCode = _fitsCodeMap[colDesc["datatype"]]
97 length = votcommon.getLength(colDesc["arraysize"])
98
99 if length==1:
100 return length, baseCode
101
102 elif length is None:
103
104
105
106
107 raise common.CannotSerializeIn("fits")
108
109 else:
110 return length, "%d%s"%(length, baseCode)
111
112
114 """returns a pyfits-capable column array for non-string values
115 stored in the colInd-th column of values.
116 """
117 length, typecode = _computeTypeCode(colDesc)
118 nullValue = _getNullValue(colDesc)
119 nan = float("NaN")
120
121 if length>1:
122
123
124
125
126
127 nullArray = numpy.array([nan]*length)
128 def mkval(v):
129 if v is None:
130 return nullArray
131 else:
132 return numpy.array(v)
133
134 else:
135 def mkval(v):
136 if v is None:
137 if nullValue is None:
138 raise ValueError("While serializing a FITS table: NULL"
139 " detected in column '%s' but no null value declared"%
140 colDesc["name"])
141 return nullValue
142 else:
143 return v
144
145 arr = numpy.array([mkval(v[colInd]) for v in values])
146 return typecode, arr
147
148
150 """returns a pyfits hdu for the valuemappers.SerManager instance table.
151 """
152 values = list(serMan.getMappedTuples())
153 columns = []
154 utypes = []
155 descriptions = []
156
157 for colInd, colDesc in enumerate(serMan):
158 descriptions.append(colDesc["description"])
159 if colDesc["datatype"]=="char" or colDesc["datatype"]=="unicodeChar":
160 makeArray = _makeStringArray
161 else:
162 makeArray = _makeValueArray
163
164 typecode, arr = makeArray(values, colInd, colDesc)
165 if typecode[-1] in 'ED':
166 nullValue = None
167 else:
168 nullValue = _getNullValue(colDesc)
169
170 columns.append(pyfits.Column(name=str(colDesc["name"]),
171 unit=str(colDesc["unit"]), format=typecode,
172 null=nullValue, array=arr))
173 if colDesc["utype"]:
174 utypes.append((colInd, str(colDesc["utype"].lower())))
175
176 hdu = pyfits.BinTableHDU.from_columns(pyfits.ColDefs(columns))
177 for colInd, utype in utypes:
178 hdu.header.set("TUTYP%d"%(colInd+1), utype)
179
180 cards = hdu.header.cards
181 for colInd, desc in enumerate(descriptions):
182 cards["TTYPE%d"%(colInd+1)].comment = desc.encode("ascii", "ignore")
183 hdu.header.set("TCOMM%d"%(colInd+1), desc.encode("ascii", "ignore"))
184
185 if not hasattr(serMan.table, "IgnoreTableParams"):
186 for param in serMan.table.iterParams():
187 if param.value is None:
188 continue
189
190 key, value, comment = str(param.name), param.value, param.description
191 if isinstance(value, unicode):
192 value = value.encode('ascii', "xmlcharrefreplace")
193 if isinstance(comment, unicode):
194 comment = comment.encode('ascii', "xmlcharrefreplace")
195 if len(key)>8:
196 key = "hierarch "+key
197
198 try:
199 hdu.header.set(key, value=value, comment=comment)
200 except ValueError as ex:
201
202 base.ui.notifyWarning(
203 "Failed to serialise param %s to a FITS header (%s)"%(
204 param.name,
205 utils.safe_str(ex)))
206
207 return hdu
208
209
211 """returns a hdulist containing extensions for the tables in dataSet.
212
213 You must make sure that this function is only executed once
214 since pyfits is not thread-safe.
215 """
216 tables = [base.SerManager(table, acquireSamples=acquireSamples)
217 for table in dataSet.tables.values()]
218 extensions = [_makeExtension(table) for table in tables]
219 primary = pyfits.PrimaryHDU()
220 primary.header.set("DATE", time.strftime("%Y-%m-%d"),
221 "Date file was written")
222 return pyfits.HDUList([primary]+extensions)
223
224
226 """returns a hdulist containing extensions for the tables in dataSet.
227
228 This function may block basically forever. Never call this from
229 the main server, always use threads or separate processes (until
230 pyfits is fixed to be thread-safe).
231
232 This will add table parameters as header cards on the resulting FITS
233 header.
234 """
235 with utils.fitsLock():
236 return _makeFITSTableNOLOCK(dataSet, acquireSamples)
237
238
240 """returns the name of a temporary file containing the FITS data for
241 hdulist.
242 """
243 handle, pathname = tempfile.mkstemp(".fits", prefix="fitstable",
244 dir=base.getConfig("tempDir"))
245
246
247
248 if len(hdulist)>1:
249 hdulist[0].header.set("EXTEND", True, "More exts following")
250
251 with utils.silence():
252 hdulist.writeto(pathname, clobber=1)
253 os.close(handle)
254 return pathname
255
256
258 """returns the name of a temporary file containing a fits file
259 representing dataSet.
260
261 The caller is responsible to remove the file.
262 """
263 hdulist = makeFITSTable(dataSet, acquireSamples)
264 return writeFITSTableFile(hdulist)
265
266
268 """a formats.common compliant data writer.
269
270 This will write out table params as header cards. To serialise
271 those yourself (as is required for spectral data model compliant
272 tables), set an attribute IgnoreTableParams (with an arbitrary
273 value) on the table.
274 """
275 data = rsc.wrapTable(data)
276 fitsName = makeFITSTableFile(data, acquireSamples)
277 try:
278 src = open(fitsName)
279 utils.cat(src, outputFile)
280 src.close()
281 finally:
282 os.unlink(fitsName)
283
284 common.registerDataWriter("fits", writeDataAsFITS, "application/fits",
285 "FITS Binary Table", ".fits")
286