Merge ad6a46a into 80ad14a

daiquiri/core/adapter/download/base.py

@@ -6,7 +6,7 @@
 
 from django.conf import settings
 
-from daiquiri.core.generators import generate_csv, generate_votable
+from daiquiri.core.generators import generate_csv, generate_votable, generate_fits
 from daiquiri.core.utils import get_doi_url
 
 logger = logging.getLogger(__name__)
@@ -18,7 +18,7 @@ def __init__(self, database_key, database_config):
         self.database_key = database_key
         self.database_config = database_config
 
-    def generate(self, format_key, schema_name, table_name, columns, sources=None, status=None, empty=False):
+    def generate(self, format_key, schema_name, table_name, columns, sources=None, status=None, nrows=None):
         # create the final list of arguments subprocess.Popen
         if format_key == 'sql':
             # create the final list of arguments subprocess.Popen
@@ -54,7 +54,12 @@ def generate(self, format_key, schema_name, table_name, columns, sources=None, s
 
                 return generate_votable(self.generate_rows(prepend=prepend), columns,
                                         resource_name=schema_name, table_name=table_name,
-                                        links=links, query_status=status, empty=empty)
+                                        links=links, query_status=status, empty=(nrows==0))
+
+            elif format_key == 'fits':
+                return generate_fits(self.generate_rows(prepend=prepend), columns,
+                                     nrows=nrows, table_name=table_name)
+
             else:
                 raise Exception('Not supported.')
 

daiquiri/core/generators.py

@@ -1,6 +1,12 @@
 import csv
+import datetime
 import io
 import sys
+import struct
+
+from django.contrib.sites.models import Site
+
+from daiquiri import __version__ as daiquiri_version
 
 
 def generate_csv(generator, fields):
@@ -99,3 +105,127 @@ def generate_votable(generator, fields, resource_name=None, table_name=None, lin
     </RESOURCE>
 </VOTABLE>
 '''
+
+
+def generate_fits(generator, fields, nrows, table_name=None):
+
+    # VO format label, FITS format label, size, NULL value, encoded value
+    formats_dict = {
+        'unsignedByte': ('s', 'L', 1,  b'\x00',             lambda x: b'T' if x == 'true' else b'F'),
+        'short':        ('h', 'I', 2,  32767,               int),
+        'int':          ('i', 'J', 4,  2147483647,          int),
+        'long':         ('q', 'K', 8,  9223372036854775807, int),
+        'float':        ('f', 'E', 4,  float('nan'),        float),
+        'double':       ('d', 'D', 8,  float('nan'),        float),
+        'char':         ('s', 'A', 32,  b'',                lambda x: x.encode()),
+        'timestamp':    ('s', 'A', 19, b'',                 lambda x: x.encode()),
+        'array':        ('s', 'A', 64, b'NULL',             lambda x: x.encode()),
+        'spoint':       ('s', 'A', 64, b'NULL',             lambda x: x.encode())
+    }
+
+    names = [d['name'] for d in fields]
+    datatypes = [d['datatype'] for d in fields]
+    arraysizes = [d['arraysize'] if d['arraysize'] is not None else ''
+                  for d in fields]
+    for i, d in enumerate(zip(datatypes, arraysizes)):
+        if d[0] == 'timestamp':
+            arraysizes[i] = formats_dict['timestamp'][2]
+        if d[0] in ('char', 'spoint', 'array') and d[1] == '':
+            arraysizes[i] = formats_dict[d[0]][2] 
+
+    naxis1 = sum([formats_dict[i[0]][2] if not i[1] else i[1]
+                  for i in zip(datatypes, arraysizes)])
+    naxis2 = nrows
+    tfields = len(names)
+
+    # Main header #############################################################
+    header0 = [i.ljust(80) for i in [
+        'SIMPLE  =                    T / conforms to FITS standard',
+        'BITPIX  =                    8 / array data type',
+        'NAXIS   =                    0 / number of array dimensions',
+        'EXTEND  =                    T',
+        'NTABLE  =                    1',
+        'END'
+        ]]
+
+    h0 = ''.join(header0)
+    h0 += ' ' * (2880 * (len(h0) // 2880 + 1) - len(h0))
+
+    yield h0.encode()
+
+    # Table header ############################################################
+    header1 = [
+       "XTENSION= 'BINTABLE'           / binary table extension".ljust(80),
+       'BITPIX  =                    8 / array data type'.ljust(80),
+       'NAXIS   =                    2 / number of array dimensions'.ljust(80),
+       'NAXIS1  = %20d / length of dimension 1'.ljust(64),
+       'NAXIS2  = %20d / length of dimension 2'.ljust(64),
+       'PCOUNT  =                    0 / number of group parameters'.ljust(80),
+       'GCOUNT  =                    1 / number of groups'.ljust(80),
+       'TFIELDS = %20d / number of table fields'.ljust(64),
+        ]
+    if table_name is not None:
+        # table_name needs to be shorter than 68 chars
+        header1.append(("EXTNAME = '%s' / table name" %
+                        str(table_name[:68])).ljust(80))
+
+    h1 = ''.join(header1) % (naxis1, naxis2, tfields)
+
+    ttype = "TTYPE%s = '%s'"
+    tform = "TFORM%s = '%s'"
+    tnull = "TNULL%s = %s"
+
+    for i, d in enumerate(zip(names, datatypes, arraysizes)):
+        temp = (ttype % (str(i + 1).ljust(2),
+                d[0][:68].ljust(8)))[:80].ljust(30)
+        temp += ' / label for column %d' % (i + 1)
+        temp = temp[:80]
+        temp += ' ' * (80 - len(temp))
+        h1 += temp
+
+        ff = (str(d[2]) + formats_dict[d[1]][1]).ljust(8)
+        temp = (tform % (str(i + 1).ljust(2), ff))[:80].ljust(31)
+        temp += '/ format for column %d' % (i + 1)
+        temp = temp[:80]
+        temp += ' ' * (80 - len(temp))
+        h1 += temp
+
+        # NULL values only for int-like types
+        if d[1] in ('short', 'int', 'long'):
+            temp = (tnull % (str(i + 1).ljust(2),
+                    formats_dict[d[1]][3]))[:80].ljust(31)
+            temp += '/ blank value for column %d' % (i + 1)
+            temp = temp[:80]
+            temp += ' ' * (80 - len(temp))
+            h1 += temp
+
+    now = datetime.datetime.utcnow().replace(microsecond=0).isoformat()
+    site = Site.objects.get_current()
+    h1 += ("DATE-HDU= '%s' / UTC date of HDU creation" % now).ljust(80)
+    h1 += ("DAIQUIRI= '%s'%s / Daiquiri version" % (daiquiri_version,
+        ' ' * max(0, 18 - len(daiquiri_version)))).ljust(80)
+    h1 += ("SOURCE  = '%s'%s / table origin" % (site,
+        ' ' * max(0, 18 - len(str(site))))).ljust(80)
+
+    h1 += 'END'.ljust(80)
+    h1 += ' ' * (2880 * (len(h1) // 2880 + 1) - len(h1))
+
+    yield h1.encode()
+
+    # Data ####################################################################
+    fmt = '>' + ''.join([str(i[1]) + formats_dict[i[0]][0]
+                         for i in zip(datatypes, arraysizes)])
+
+    row_count = 0
+    for row in generator:
+        r = [formats_dict[i[1]][3] if i[0] == 'NULL'
+             else formats_dict[i[1]][4](i[0]) for i in zip(row, datatypes)]
+
+        yield struct.pack(fmt, *r)
+        row_count += 1
+
+    # Footer padding (to fill the last block to 2880 bytes) ###################
+    ln = naxis1 * row_count
+    footer = '\x00' * (2880 * (ln // 2880 + 1) - ln)
+
+    yield footer.encode()

daiquiri/core/settings/daiquiri.py

@@ -141,6 +141,13 @@
         'content_type': 'application/xml',
         'label': 'IVOA VOTable XML file - TABLEDATA serialization',
         'help': 'A XML file using the IVOA VOTable format. Use this option if you intend to use VO compatible software to further process the data.'
+    },
+    {
+        'key': 'fits',
+        'extension': 'fits',
+        'content_type': 'application/fits',
+        'label': 'FITS',
+        'help': 'Flexible Image Transport System (FITS) file format.'
     }
 ]
 

daiquiri/query/models.py

@@ -366,7 +366,7 @@ def stream(self, format_key):
                 self.metadata.get('columns'),
                 self.metadata.get('sources'),
                 self.result_status,
-                (self.nrows == 0)
+                self.nrows
             )
         else:
             raise ValidationError({

daiquiri/query/tasks.py

@@ -232,8 +232,13 @@ def create_download_file(download_id):
         download_job.save()
 
         # write file using the generator in the adapter
+        if download_job.format_key == 'fits':
+            write_label = 'wb'
+        else:
+            write_label = 'w'
         try:
-            with open(download_job.file_path, 'w') as f:
+
+            with open(download_job.file_path, write_label) as f:
                 for line in download_job.job.stream(download_job.format_key):
                     f.write(line)