Add packets_to_rdrs function

edosl0util/rdrgen.py

@@ -3,6 +3,7 @@ import itertools
 import os
 from collections import OrderedDict
 from datetime import datetime
+from tempfile import TemporaryFile
 
 import attr
 import h5py
@@ -12,6 +13,29 @@ from astropy.time import Time, TimeDelta
 import edosl0util
 from edosl0util.jpssrdr import (
     StaticHeader, Apid as ApidListItem, PacketTracker, decode_rdr_blob)
+from edosl0util.stream import jpss_packet_stream
+
+
+def packets_to_rdrs(sat, pkt_files):
+    # FIXME: refactor!!!
+    rdr_pkt_files = {}
+    for pkt_file in pkt_files:
+        with open(pkt_file) as file_obj:
+            stream = jpss_packet_stream(file_obj)
+            for pkt in stream:
+                rdr_type = get_rdr_type(pkt.apid)
+                gran = calc_rdr_granule(sat, rdr_type, pkt)
+                if (rdr_type, gran) not in rdr_pkt_files:
+                    rdr_pkt_files[rdr_type, gran] = TemporaryFile()
+                rdr_pkt_files[rdr_type, gran].write(pkt.bytes())
+    for rdr_pkt_file in rdr_pkt_files.values():
+        rdr_pkt_file.seek(0)
+        pkts = jpss_packet_stream(rdr_pkt_file)
+        pkts = sorted(pkts, key=(lambda p: (timecode_to_iet(p.secondary_header.timecode),
+                                            p.apid)))
+        blob = build_rdr_blob(sat, pkts)
+        write_rdr(sat, blob)
+
 
 
 def write_rdr(sat, blob, dir_path='.',
@@ -26,7 +50,7 @@ def write_rdr(sat, blob, dir_path='.',
     software_ver = (software_ver
                     or edosl0util.__name__ + '-' + edosl0util.__version__)
     blob_info = decode_rdr_blob(blob)
-    rdr_type = SpacecraftDiaryRdrType()  # TODO: select via APID
+    rdr_type = get_rdr_type(blob_info.apids[0].value)
     gran_iet = blob_info.header.start_boundary
     gran_end_iet = blob_info.header.end_boundary
     gran_time = iet_to_datetime(gran_iet)
@@ -108,10 +132,10 @@ def set_h5_attrs(h5_obj, attrs):
 
 
 def build_rdr_blob(sat, pkt_stream):
+    pkt_stream = iter(pkt_stream)
     first_pkt = next(pkt_stream)  # FIXME: what if there are no packets?
-    rdr_type = SpacecraftDiaryRdrType()  # TODO: select based on first packet APID
-    granule_iet = calc_iet_granule(satellite_base_times[sat], rdr_type.gran_len,
-                                   timecode_to_iet(first_pkt.secondary_header.timecode))
+    rdr_type = get_rdr_type(first_pkt.apid)
+    granule_iet = calc_rdr_granule(sat, rdr_type, first_pkt)
     granule_iet_end = granule_iet + rdr_type.gran_len
 
     total_pkt_size = 0
@@ -230,6 +254,27 @@ class ApidSpec(object):
     max_expected = attr.ib()
 
 
+class RdrTypeManager(object):
+    def __init__(self):
+        self._types_by_apid = {}
+
+    def register_type(self, cls):
+        for apid_spec in cls.apids:
+            if apid_spec.num in self._types_by_apid:
+                raise ValueError('each APID can only be handled by one RDR type')
+            self._types_by_apid[apid_spec.num] = cls
+        return cls
+
+    def get_type_for_apid(self, apid):
+        return self._types_by_apid[apid]
+
+
+rdr_type_mgr = RdrTypeManager()
+rdr_type_spec = rdr_type_mgr.register_type
+get_rdr_type = rdr_type_mgr.get_type_for_apid
+
+
+@rdr_type_spec
 class CrisScienceRdrType(object):
     product_id = 'RCRIS'
     short_name = 'CRIS-SCIENCE-RDR'
@@ -240,6 +285,7 @@ class CrisScienceRdrType(object):
     apids = get_cris_science_apids()
 
 
+@rdr_type_spec
 class SpacecraftDiaryRdrType(object):
     product_id = 'RNSCA'
     short_name = 'SPACECRAFT-DIARY-RDR'
@@ -262,7 +308,7 @@ def make_rdr_filename(prod_id, sat, gran_time, gran_end_time, orbit_num, creatio
             return o[:-1] + ('c' if compressed else 'u')
         else:
             return o
-    return '{p}_{s}_d{d:%Y%m%d}_t{b}_e{e}_b{n}_c{c:%Y%m%d%H%M%S%f}_{o}_{m}.h5'.format(
+    return '{p}_{s}_d{d:%Y%m%d}_t{b}_e{e}_b{n:05d}_c{c:%Y%m%d%H%M%S%f}_{o}_{m}.h5'.format(
         p=prod_id, s=sat_strs[sat], d=gran_time, b=format_gran_time(gran_time),
         e=format_gran_time(gran_end_time), n=orbit_num, c=creation_time,
         o=format_origin(origin), m=domain)
@@ -279,6 +325,11 @@ def iet_to_datetime(iet):
     return (iet_epoch + TimeDelta(iet * 1e-6, format='sec')).utc.datetime
 
 
+def calc_rdr_granule(sat, rdr_type, pkt):
+    return calc_iet_granule(satellite_base_times[sat], rdr_type.gran_len,
+                            timecode_to_iet(pkt.secondary_header.timecode))
+
+
 def calc_iet_granule(base_time, gran_len, iet):
     return (iet - base_time) // gran_len * gran_len + base_time