speed up is_compatible; makes event_source factory faster

ctapipe/io/lsteventsource.py

@@ -25,10 +25,10 @@ class LSTEventSource(EventSource):
     EventSource for LST r0 data.
     """
 
-    
+
 
     def __init__(self, config=None, tool=None, **kwargs):
-        
+
         """
         Constructor
         Parameters
@@ -63,7 +63,7 @@ def __init__(self, config=None, tool=None, **kwargs):
 
 
         self.multi_file = MultiFiles(self.file_list)
-        
+
         self.camera_config = self.multi_file.camera_config
         self.log.info("Read {} input files".format(self.multi_file.num_inputs()))
 
@@ -123,6 +123,10 @@ def _generator(self):
 
     @staticmethod
     def is_compatible(file_path):
+        from .sst1meventsource import is_fits_in_header
+        if not is_fits_in_header(file_path):
+            return False
+
         from astropy.io import fits
         try:
             # The file contains two tables:
@@ -215,8 +219,8 @@ def fill_r0_camera_container_from_zfile(self, container, event):
 
         reshaped_waveform = np.array(
                 event.waveform
-             ).reshape(n_gains, 
-                       self.camera_config.num_pixels, 
+             ).reshape(n_gains,
+                       self.camera_config.num_pixels,
                        container.num_samples)
 
         # initialize the waveform container to zero
@@ -247,7 +251,7 @@ class MultiFiles:
 
     """
     This class open all the files in file_list and read the events following
-    the event_id order 
+    the event_id order
     """
 
     def __init__(self, file_list):
@@ -258,9 +262,9 @@ def __init__(self, file_list):
         self._camera_config = {}
         self.camera_config = None
 
-    
+
         paths = []
-        for file_name in file_list:            
+        for file_name in file_list:
             paths.append(file_name)
             Provenance().add_input_file(file_name, role='r0.sub.evt')
 

ctapipe/io/nectarcameventsource.py

@@ -23,8 +23,8 @@ class NectarCAMEventSource(EventSource):
     """
 
     def __init__(self, config=None, tool=None, **kwargs):
-   
-        
+
+
         """
         Constructor
         Parameters
@@ -87,7 +87,7 @@ def _generator(self):
             tel_descr = TelescopeDescription(optics, camera)
 
             tel_descr.optics.tel_subtype = ''  # to correct bug in reading
-            
+
             self.n_camera_pixels = tel_descr.camera.n_pixels
             tels = {tel_id: tel_descr}
 
@@ -118,6 +118,10 @@ def _generator(self):
 
     @staticmethod
     def is_compatible(file_path):
+        from .sst1meventsource import is_fits_in_header
+        if not is_fits_in_header(file_path):
+            return False
+
         from astropy.io import fits
         try:
             # The file contains two tables:
@@ -211,8 +215,8 @@ def fill_r0_camera_container_from_zfile(self, container, event):
 
         reshaped_waveform = np.array(
             event.waveform
-             ).reshape(n_gains, 
-                       self.camera_config.num_pixels, 
+             ).reshape(n_gains,
+                       self.camera_config.num_pixels,
                        container.num_samples)
 
         # initialize the waveform container to zero

ctapipe/io/sst1meventsource.py

@@ -4,6 +4,7 @@
 
 Needs protozfits v1.0.2 from github.com/cta-sst-1m/protozfitsreader
 """
+import gzip
 import numpy as np
 from .eventsource import EventSource
 from .containers import SST1MDataContainer
@@ -12,6 +13,24 @@
 __all__ = ['SST1MEventSource']
 
 
+def is_fits_in_header(file_path):
+    '''quick check if file is a FITS file
+
+    by looking into the first 1024 bytes and searching for the string "FITS"
+    typically used in is_compatible
+    '''
+    # read the first 1kB
+    with open(file_path, 'rb') as f:
+        marker_bytes = f.read(1024)
+
+    # if file is gzip, read the first 4 bytes with gzip again
+    if marker_bytes[0] == 0x1f and marker_bytes[1] == 0x8b:
+        with gzip.open(file_path, 'rb') as f:
+            marker_bytes = f.read(1024)
+
+    return b'FITS' in marker_bytes
+
+
 class SST1MEventSource(EventSource):
 
     def __init__(self, config=None, tool=None, **kwargs):
@@ -78,6 +97,9 @@ def _generator(self):
 
     @staticmethod
     def is_compatible(file_path):
+        if not is_fits_in_header(file_path):
+            return False
+
         from astropy.io import fits
         try:
             h = fits.open(file_path)[1].header