Fix bug in apid list. Cache pkt structs by size.

edosl0util/jpssrdr.py

@@ -64,7 +64,7 @@ class PacketTracker(BaseStruct):
     ]
 
 
-def _make_packet(buf, size, offset):
+def _make_packet_impl(size):
     """
     Create a struct for a CCSDS packet. The struct size is dynamic so we need
     a new class for each packet.
@@ -85,9 +85,7 @@ def _make_packet(buf, size, offset):
             ('length_minus1', c.c_uint16),
             ('data', c.c_uint8 * data_size)
         ]
-    pkt = PktImpl.from_buffer(buf, offset)
-    assert pkt.length_minus1 + 1 == data_size
-    return pkt
+    return PktImpl
 
 
 Packet = namedtuple('Packet', ('tracker', 'packet'))
@@ -113,6 +111,10 @@ class CommonRdr(namedtuple('CommonRdr', ('buf', 'header', 'apids'))):
 
 
 def _packets_for_apid(buf, header, apid):
+
+    # we cache packet classes by their size
+    _pkt_impl_cache = {}
+
     t_off = header.pkt_tracker_offset + apid.pkt_tracker_start_idx * c.sizeof(PacketTracker)
     for idx in range(apid.pkts_received):
         tracker = PacketTracker.from_buffer(buf, t_off)
@@ -121,18 +123,24 @@ def _packets_for_apid(buf, header, apid):
             continue
 
         p_off = header.ap_storage_offset + tracker.offset
-        # packet = buf[p_off:p_off+tracker.size]
-        yield tracker, _make_packet(buf, tracker.size, p_off)
+        pkt_impl = _pkt_impl_cache.get(tracker.size)
+        if not pkt_impl:
+            pkt_impl = _make_packet_impl(tracker.size)
+            _pkt_impl_cache[tracker.size] = pkt_impl
+
+        pkt = pkt_impl.from_buffer(buf, p_off)
+        assert c.sizeof(pkt) == tracker.size
+        yield tracker, pkt
 
 
 def _read_apid_list(header, buf):
     """
     Return a generator that yields `Apid`s
     """
-    start = header.apid_list_offset
-    end = header.num_apids * c.sizeof(Apid)
-    for offset in range(start, end, c.sizeof(Apid)):
+    offset = header.apid_list_offset
+    for idx in range(header.num_apids):
         yield Apid.from_buffer(buf, offset)
+        offset += c.sizeof(Apid)
 
 
 def read_common_rdrs(sensor, filepath):
@@ -160,7 +168,7 @@ def read_rdr_datasets(sensor, filepath):
     cmp_names = lambda *tup: cmp(*(int(x.split('_')[-1]) for x in tup))  # noqa
     for name in sorted(dsnames, cmp=cmp_names):
         ds = grp[name]
-        yield np.getbuffer(np.array(ds))
+        yield np.array(ds)
 
 
 def sort_packets_by_obs_time(packets):
@@ -192,31 +200,46 @@ def sort_packets_edos(rdr, packets):
     return sort_packets_by_obs_time(sort_packets_by_apid(packets, order=order))
 
 
-def convert_to_nasa_l0(sensor, filename, skipfill=False, outfn=None):
+def convert_to_nasa_l0(sensor, filename, skipfill=False):
     """
     Convert a JSPP RDR to a NASA L0 PDS file.
     """
-    outfn = outfn or '{}.pds'.format(filename)
-    with open(outfn, 'wb') as fptr:
-        for idx, rdr in enumerate(read_common_rdrs(sensor, filename)):
-            packets = sort_packets_edos(rdr, rdr.packets())
-            for packet in packets:
-                if packet.tracker.fill_percent != 0:
-                    LOG.debug('fill: %s', packet.tracker)
-                    if skipfill:
-                        continue
-                fptr.write(packet.packet)
-    return outfn
+    for idx, rdr in enumerate(read_common_rdrs(sensor, filename)):
+        LOG.debug(rdr.header)
+        for apid in rdr.apids:
+            LOG.debug(apid)
+        packets = sort_packets_edos(rdr, rdr.packets())
+        for packet in packets:
+            if packet.tracker.fill_percent != 0:
+                LOG.debug('fill: %s', packet.tracker)
+                if skipfill:
+                    continue
+            yield packet.packet
 
 
 if __name__ == '__main__':
     import argparse
+    from datetime import datetime
     parser = argparse.ArgumentParser()
+    parser.add_argument('-v', action='store_true')
     parser.add_argument('-o', '--output')
     parser.add_argument('-f', '--skipfill', action='store_true')
+    def interval(val):
+        fmt = '%Y-%m-%d %H:%M:%S'
+        start, end = val.split(',')
+        return datetime.strptime(start, fmt), datetime.strptime(end, fmt)
+    parser.add_argument('-t', '--trunc', type=interval)
     parser.add_argument('sensor', choices=('viirs', 'atms', 'cris'))
     parser.add_argument('rdr')
     args = parser.parse_args()
 
-    logging.basicConfig(level=logging.DEBUG)
-    convert_to_nasa_l0(args.sensor, args.rdr, outfn=args.output)
+    if args.v:
+        lvl = logging.DEBUG
+    else:
+        lvl = logging.WARN
+    logging.basicConfig(level=lvl, format='%(message)s')
+
+    output = args.output or args.rdr + '.pds'
+    with open(output, 'wb') as fptr:
+        for packet in convert_to_nasa_l0(args.sensor, args.rdr):
+            fptr.write(packet)