NASA Ocean Color

ocssw V2022

 """Utils for the CCSDSPy package."""
  
 __author__ = "Daniel da Silva <mail@danieldasilva.org>"
  
 from io import BytesIO
 import warnings
  
 import numpy as np
  
 from . import VariableLength, PacketArray
 from .constants import BITS_PER_BYTE, PRIMARY_HEADER_NUM_BYTES
 from . import decode
  
  
 def get_packet_total_bytes(primary_header_bytes):
     """[autogenerated]"""
     # Wrap around internal function from decode module
     return decode._get_packet_total_bytes(primary_header_bytes)
  
  
 def get_packet_apid(primary_header_bytes):
     """[autogenerated]"""
     # Wrap around internal function from decode module
     return decode._get_packet_apid(primary_header_bytes)
  
  
 # Copy __doc__ from decode module for functions which wrap decode module functions
 # so that the doc need not be repeated in two places.
 get_packet_total_bytes.__doc__ = decode._get_packet_total_bytes.__doc__
 get_packet_apid.__doc__ = decode._get_packet_apid.__doc__
  
  
 def iter_packet_bytes(file, include_primary_header=True):
     """Iterate through packets as raw bytes objects, in the order they appear in a file.
  
     This function works with mixed files containing multiple APIDs, which may
     include both fixed length and variable length packets.
  
     If end of last packet doesn't align with end of file, a warning is issued.
  
     Parameters
     ----------
     file : str, file-like
       Path to file on the local file system, or file-like object
     include_primary_header : bool
       If set to False, excludes the primary header bytes (the first six)
  
     Yields
     ------
     packet_bytes : bytes
        Bytes associated with each packet as it appears in the file. When
        include_primary_header=False, the primary header bytes are excluded.
     """
     if hasattr(file, "read"):
         file_bytes = np.frombuffer(file.read(), "u1")
     else:
         file_bytes = np.fromfile(file, "u1")
  
     offset = 0
  
     if include_primary_header:
         delta_idx = 0
     else:
         delta_idx = PRIMARY_HEADER_NUM_BYTES
  
     while offset < len(file_bytes):
         packet_nbytes = get_packet_total_bytes(
             file_bytes[offset : offset + PRIMARY_HEADER_NUM_BYTES].tobytes()
         )
         packet_bytes = file_bytes[offset + delta_idx : offset + packet_nbytes].tobytes()
  
         yield packet_bytes
  
         offset += packet_nbytes
  
     if offset != len(file_bytes):
         missing_bytes = offset - len(file_bytes)
         message = (
             f"File appears truncated-- missing {missing_bytes} byte (or " "maybe garbage at end)"
         )
         warnings.warn(message)
  
  
 def split_packet_bytes(file, include_primary_header=True):
     """Retrieve a list of bytes objects corresponding to each packet in a file.
  
     This function works with mixed files containing multiple APIDs, which may
     include both fixed length and variable length packets.
  
     If end of last packet doesn't align with end of file, a warning is issued.
  
     Parameters
     ----------
     file : str, file-like
       Path to file on the local file system, or file-like object
     include_primary_header : bool
       If set to False, excludes the primary header bytes (the first six)
  
     Returns
     -------
     packet_bytes : list of bytes
        List of bytes objects associated with each packet as it appears in the
        file. When include_primary_header=False, each byte object will have its
        primary header bytes excluded.
     """
     return list(iter_packet_bytes(file, include_primary_header=include_primary_header))
  
  
 def read_primary_headers(file):
     """Read primary header fields and return contents as a dictionary
     of arrays.
  
     This function works with mixed files containing multiple APIDs, which may
     include both fixed length and variable length packets.
  
     Parameters
     ----------
     file : str, file-like
       Path to file on the local file system, or file-like object
  
     Returns
     -------
     header_arrays : dict, string to NumPy array
        Dictionary mapping header names to NumPy arrays. The header names are:
        `CCSDS_VERSION_NUMBER`, `CCSDS_PACKET_TYPE`, `CCSDS_SECONDARY_FLAG`,
        `CCSDS_SEQUENCE_FLAG`, `CCSDS_APID`, `CCSDS_SEQUENCE_COUNT`,
        `CCSDS_PACKET_LENGTH`
     """
     pkt = VariableLength(
         [
             PacketArray(
                 name="unused", data_type="uint", bit_length=BITS_PER_BYTE, array_shape="expand"
             )
         ]
     )
  
     header_arrays = pkt.load(file, include_primary_header=True)
     del header_arrays["unused"]
  
     return header_arrays
  
  
 def split_by_apid(mixed_file, valid_apids=None):
     """Split a stream of mixed APIDs into separate streams by APID.
  
     This works with a mix of both fixed length and variable length packets.
  
     Parameters
     ----------
     mixed_file: str, file-like
        Path to file on the local file system, or file-like object
     valid_apids: list of int, None
        Optional list of valid APIDs. If specified, warning will be issued when
        an APID is encountered outside this list.
  
     Returns
     -------
     stream_by_apid : dict, int to :py:class:`~io.BytesIO`
       Dictionary mapping integer apid number to BytesIO instance with the file
       pointer at the beginning of the stream.
     """
     # If not None, convert valid_apids to set for faster lookup times
     if valid_apids is not None:
         valid_apids = set(valid_apids)
  
     stream_by_apid = {}
  
     for packet_bytes in iter_packet_bytes(mixed_file):
         apid = get_packet_apid(packet_bytes[:PRIMARY_HEADER_NUM_BYTES])
  
         if valid_apids is not None and apid not in valid_apids:
             warnings.warn(f"Found unknown APID {apid}")
  
         if apid not in stream_by_apid:
             stream_by_apid[apid] = BytesIO()
  
         stream_by_apid[apid].write(packet_bytes)
  
     for stream in stream_by_apid.values():
         stream.seek(0)
  
     return stream_by_apid
  
  
 def count_packets(file, return_missing_bytes=False):
     """Count the number of packets in a file and check if there are any
     missing bytes in the last packet.
  
     This function works with mixed files containing multiple APIDs, which may
     include both fixed length and variable length packets. When used with
     multiple APIDs, it simply returns the total number of packets of any APID.
  
     If end of last packet doesn't align with end of file, a warning is issued.
  
     Parameters
     ----------
     file : str, file-like
       Path to file on the local file system, or file-like object
     return_missing_bytes : bool, optional
       Also return the number of missing bytes at the end of the file. This
       is the number of bytes which would need to be added to the file to
       complete the last packet expected (as set by the packet length in
       the last packet's primary header).
  
     Returns
     -------
     num_packets : int
        Number of complete packets in the file
     missing_bytes : int, optional
       The number of bytes which would need to be added to the file to
       complete the last packet expected (as set by the packet length in
       the last packet's primary header).
     """
     if hasattr(file, "read"):
         file_bytes = np.frombuffer(file.read(), "u1")
     else:
         file_bytes = np.fromfile(file, "u1")
  
     offset = 0
     num_packets = 0
  
     while offset < len(file_bytes):
         packet_nbytes = get_packet_total_bytes(
             file_bytes[offset : offset + PRIMARY_HEADER_NUM_BYTES].tobytes()
         )
         offset += packet_nbytes
         num_packets += 1
  
     missing_bytes = offset - len(file_bytes)
  
     if offset != len(file_bytes):
         missing_bytes = offset - len(file_bytes)
         message = (
             f"File appears truncated-- missing {missing_bytes} byte (or " "maybe garbage at end)"
         )
         warnings.warn(message)
  
     if return_missing_bytes:
         return num_packets, missing_bytes
     else:
         return num_packets