Skymap is passed around as HDUList. This implements #10

tests/test_torosgcn.py

@@ -549,6 +549,27 @@ def test_generate_targets(self, mock_config):
 
         now = Time("2019-06-01T01:00:00")
         obs_ret = torosgcn.scheduler.generate_targets(skymap_path, detection_time=now)
+        self.assertEqual(len(obs_ret), 2)
+        eaba = obs_ret[0]
+        ctmo = obs_ret[1]
+        self.assertTrue("name" in eaba)
+        self.assertEqual("EABA", eaba.get("name"))
+        self.assertTrue("name" in ctmo)
+        self.assertEqual("CTMO", ctmo.get("name"))
+        self.assertTrue("targets" in eaba)
+        for targets in [eaba.get("targets"), ctmo.get("targets")]:
+            self.assertTrue("Name" in targets.colnames)
+            self.assertTrue("RA" in targets.colnames)
+            self.assertTrue("Dec" in targets.colnames)
+            self.assertTrue("Likelihood" in targets.colnames)
+            self.assertEqual(len(targets), 5)
+
+        # Test that it works if skymap is an hdulist
+        from astropy.io import fits
+        hdulist = fits.open(skymap_path)
+        fp.seek(0)
+        mock_config.get_config_for_key.side_effect = [self.obs, fp.name, cat_filters]
+        obs_ret = torosgcn.scheduler.generate_targets(hdulist, detection_time=now)
         fp.close()
         self.assertEqual(len(obs_ret), 2)
         eaba = obs_ret[0]

torosgcn/listen.py

@@ -206,19 +206,21 @@ def getinfo(root):
 
 
 def retrieve_skymap(info):
-    "Fetch skymap FITS file from GraceDB service. Return skymap FITS as raw bytes"
+    "Fetch skymap FITS file from GraceDB service. Return skymap FITS as HDUList"
     import requests
     import tempfile
+    from astropy.io import fits
 
     fits_response = requests.get(info.get("skymap_fits"), stream=False)
-    fits_response.raise_for_status()
+    if not fits_response.ok:
+        logger.error("Downloading Skymap response not OK.")
     fp = tempfile.NamedTemporaryFile()
     for block in fits_response.iter_content(1024):
         fp.write(block)
     fp.seek(0)
-    skymap_data = fp.read()
+    skymap_hdulist = fits.open(fp.name)
     fp.close()
-    return skymap_data
+    return skymap_hdulist
 
 
 def upload_gcnnotice(info, targets=None):
@@ -260,7 +262,7 @@ def upload_gcnnotice(info, targets=None):
     logger.debug("Target JSON sent to broker:\n{}".format(targetsjson))
 
 
-def backup_skymap(skymap_data, info):
+def backup_skymap(skymap_hdulist, info):
     "Save to file the skymap FITS file asociated with this GCN Notice"
     bkp_config = config.get_config_for_key("Backup")
     if not bkp_config.get("Backup Skymap"):
@@ -274,8 +276,7 @@ def backup_skymap(skymap_data, info):
     )
 
     skymap_path = os.path.join(skymap_bkpdir, skymap_bkpname)
-    with open(skymap_path, "wb") as fp:
-        fp.write(skymap_data)
+    skymap_hdulist.writeto(skymap_path)
     logger.info("Skymap file {} copied to {}".format(skymap_bkpname, skymap_bkpdir))
 
 
@@ -343,13 +344,13 @@ def process_gcn(payload, root):
     targets = None
     if info.get("skymap_fits") is not None:
         try:
-            skymap_data = retrieve_skymap(info)
+            skymap_hdulist = retrieve_skymap(info)
             try:
-                backup_skymap(skymap_data, info)
+                backup_skymap(skymap_hdulist, info)
             except:
                 logger.exception("Problem backing-up skymap")
             try:
-                targets = scheduler.generate_targets(skymap_data)
+                targets = scheduler.generate_targets(skymap_hdulist)
             except:
                 logger.exception("Error generating targets")
         except:

torosgcn/scheduler.py

@@ -79,29 +79,13 @@ def generate_targets(skymap, detection_time=None):
     import healpy as hp
     from scipy.stats import norm
 
-    skymap_path = None
-    if isinstance(skymap, bytes):
-        from tempfile import NamedTemporaryFile
-
-        fp = NamedTemporaryFile()
-        skymap_path = fp.name
-        fp.write(skymap)
-        fp.seek(0)
-    elif isinstance(skymap, str):
-        skymap_path = skymap
-    else:
-        raise TypeError(
-            "{} for first argument should be 'bytes', or 'str' with file path.".format(
-                type(skymap)
-            )
-        )
     contains_dist_estimation = True
     try:
         aligo_banana, distmu, distsigma, distnorm = hp.read_map(
-            skymap_path, verbose=False, field=range(4)
+            skymap, verbose=False, field=range(4)
         )
     except:
-        aligo_banana = hp.read_map(skymap_path, verbose=False)
+        aligo_banana = hp.read_map(skymap, verbose=False)
         contains_dist_estimation = False
 
     npix = len(aligo_banana)