From 586a5c97e7481787a1f7f729a82b8d89c920623e Mon Sep 17 00:00:00 2001
From: Marten van Kerkwijk <mhvk@astro.utoronto.ca>
Date: Fri, 31 Jul 2020 15:38:48 -0400
Subject: [PATCH] Merge pull request #453 from mhvk/nuppi-tentative-support

Assume channels first rather than erroring on unknown PKTFMT.
---
 CHANGES.rst                        |  3 +++
 baseband/guppi/file_info.py        | 10 ++++++++++
 baseband/guppi/header.py           | 21 ++++++++++++++++-----
 baseband/guppi/tests/test_guppi.py | 16 ++++++++++++++++
 setup.cfg                          |  1 +
 5 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 8f79ee97..3a52c73d 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -4,6 +4,9 @@
 Bug Fixes
 ---------
 
+- Allow the GUPPI reader to assume channel-first ordering by default, i.e.,
+  no longer insist that PKTFMT is one of '1SFA' or 'SIMPLE'. Instead, ``info``
+  will include a warning for formats not known to work. [#453]
 
 4.0 (2020-07-18)
 ================
diff --git a/baseband/guppi/file_info.py b/baseband/guppi/file_info.py
index 4da5ba14..39afb606 100644
--- a/baseband/guppi/file_info.py
+++ b/baseband/guppi/file_info.py
@@ -13,6 +13,7 @@ class GUPPIFileReaderInfo(FileReaderInfo):
     # Get sample_rate from header rather than calculate it from frame_rate
     # and samples_per_frame, since we need to correct for overlap.
     attr_names = list(FileReaderInfo.attr_names)
+    attr_names.insert(attr_names.index('format')+1, 'pktfmt')
     attr_names.insert(attr_names.index('samples_per_frame')+1, 'overlap')
     attr_names = tuple(attr_names)
     """Attributes that the container provides."""
@@ -21,3 +22,12 @@ class GUPPIFileReaderInfo(FileReaderInfo):
         'Number of complete samples that overlap between frames.'))
     sample_rate = info_item(needs='header0', doc=(
         'Number of complete samples per second.'))
+
+    @info_item(needs='header0')
+    def pktfmt(self):
+        """Packet format for the data."""
+        pktfmt = self.header0['PKTFMT']
+        if pktfmt not in self.header0.supported_formats:
+            self.warnings['pktfmt'] = (f'Unknown pktfmt {pktfmt!r}. '
+                                       f'Assuming channels are stored first.')
+        return pktfmt
diff --git a/baseband/guppi/header.py b/baseband/guppi/header.py
index dbc710fc..33710087 100644
--- a/baseband/guppi/header.py
+++ b/baseband/guppi/header.py
@@ -67,6 +67,18 @@ class GUPPIHeader(fits.Header):
                  ('NPOL', 1),
                  ('OBSNCHAN', 1)]
 
+    supported_formats = {'1SFA', 'SIMPLE'}
+    """GUPPI formats that are known to work.
+
+    '1SFA' is used for all modes other than FAST4K (which is only used
+    for total intensity).  'SIMPLE' is from DSPSR, and used to support
+    time-first payloads.  See
+    https://safe.nrao.edu/wiki/pub/Main/JoeBrandt/guppi_status_shmem.pdf
+
+    If a format is not in this set, yet is known to work, a PR would be
+    most welcome.
+    """
+
     def __init__(self, *args, verify=True, mutable=True, **kwargs):
         # Comments handled by fits.Header__init__().
         super().__init__(*args, **kwargs)
@@ -77,13 +89,12 @@ def __init__(self, *args, verify=True, mutable=True, **kwargs):
 
     def verify(self):
         """Basic check of integrity."""
+        # Same check as dspsr's dsp::GUPPIFile::is_valid
         assert all(key in self for key in ('BLOCSIZE',
                                            'PKTIDX'))
-        # '1SFA' is used for all modes other than FAST4K (which is only used
-        # for total intensity).  'SIMPLE' is from DSPSR, and used to support
-        # time-first payloads.  See
-        # https://safe.nrao.edu/wiki/pub/Main/JoeBrandt/guppi_status_shmem.pdf
-        assert self['PKTFMT'] in ('1SFA', 'SIMPLE')
+        # We could check here for self['PKTFMT'] in self.supported_formats
+        # but that would break reading of unsupported but working formats,
+        # so instead this becomes just a warning in file_info.
 
     def copy(self):
         """Create a mutable and independent copy of the header."""
diff --git a/baseband/guppi/tests/test_guppi.py b/baseband/guppi/tests/test_guppi.py
index 3eb043ea..41364bdf 100644
--- a/baseband/guppi/tests/test_guppi.py
+++ b/baseband/guppi/tests/test_guppi.py
@@ -304,6 +304,22 @@ def test_file_info(self):
             assert info.frame_rate == header.sample_rate / (
                 header.samples_per_frame - header.overlap)
 
+    def test_file_info_unsupported_format(self, tmpdir):
+        filename = str(tmpdir.join('file.uppi'))
+        with guppi.open(SAMPLE_FILE, 'rb') as fh:
+            f = fh.read_frame()
+            f.header = f.header.copy()
+            f['PKTFMT'] = 'unknown'
+            with guppi.open(filename, 'wb') as fw:
+                fw.write_frame(f)
+
+        with guppi.open(filename, 'rb') as fr:
+            info = fr.info
+
+        assert info.pktfmt == 'unknown'
+        assert 'pktfmt' in info.warnings
+        assert 'Unknown pktfmt' in info.warnings['pktfmt']
+
     def test_frame(self, tmpdir):
         with guppi.open(SAMPLE_FILE, 'rb') as fh:
             frame = fh.read_frame(memmap=False)
diff --git a/setup.cfg b/setup.cfg
index 32b63d9c..a335c049 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -35,6 +35,7 @@ baseband.io =
 all =
     baseband-tasks[all]
 test =
+    pytest<6.0
     pytest-astropy-header
     pytest-doctestplus
     pytest-openfiles