unify ApiCatalog with DataProduct

dsavchenko · dsavchenko · commit 96d238367619 · 2026-04-14T19:00:14.000+02:00
diff --git a/oda_api/data_products.py b/oda_api/data_products.py
@@ -733,61 +733,100 @@ def from_file(
     def write_file(self, file_path, overwrite=True):
         self.write_fits_file(file_path, overwrite=overwrite)
 
-class ApiCatalog(object):
-
-
-    def __init__(self,cat_dict,name=None):
+class ApiCatalog(DataProduct):
+    def __init__(self, catalog_data: dict | Table, name: str | None = None):
         self.name=name
         _skip_list=['meta_ID']
-        meta = {}
         
-        lon_name = None
-        if 'cat_lon_name' in cat_dict.keys():
-            lon_name =  cat_dict['cat_lon_name']
+        if isinstance(catalog_data, dict):
+            meta = {}
 
-        lat_name = None
-        if 'cat_lat_name' in cat_dict.keys():
-            lat_name = cat_dict['cat_lat_name']
+            lon_name = None
+            if 'cat_lon_name' in catalog_data.keys():
+                lon_name =  catalog_data['cat_lon_name']
 
-        frame = None
-        if 'cat_frame' in cat_dict.keys():
-            frame = cat_dict['cat_frame']
+            lat_name = None
+            if 'cat_lat_name' in catalog_data.keys():
+                lat_name = catalog_data['cat_lat_name']
 
-        coord_units = None
-        if 'cat_coord_units' in cat_dict.keys():
-            coord_units = cat_dict['cat_coord_units']
+            frame = None
+            if 'cat_frame' in catalog_data.keys():
+                frame = catalog_data['cat_frame']
 
-        if 'cat_meta' in cat_dict.keys():
-            cat_meta_entry = cat_dict['cat_meta']
-            meta.update(cat_meta_entry)
-        
-        meta['FRAME'] = frame
-        meta['COORD_UNIT'] = coord_units
-        meta['LON_NAME'] = lon_name
-        meta['LAT_NAME'] = lat_name
+            coord_units = None
+            if 'cat_coord_units' in catalog_data.keys():
+                coord_units = catalog_data['cat_coord_units']
+
+            if 'cat_meta' in catalog_data.keys():
+                cat_meta_entry = catalog_data['cat_meta']
+                meta.update(cat_meta_entry)
+            
+            meta['FRAME'] = frame
+            meta['COORD_UNIT'] = coord_units
+            meta['LON_NAME'] = lon_name
+            meta['LAT_NAME'] = lat_name
+
+            self.table =Table(catalog_data['cat_column_list'], names=catalog_data['cat_column_names'],meta=meta)
 
-        self.table =Table(cat_dict['cat_column_list'], names=cat_dict['cat_column_names'],meta=meta)
+            if coord_units is not None:
+                self.table[lon_name]=Angle(self.table[lon_name],unit=coord_units)
+                self.table[lat_name]=Angle(self.table[lat_name],unit=coord_units)
 
-        if coord_units is not None:
-            self.table[lon_name]=Angle(self.table[lon_name],unit=coord_units)
-            self.table[lat_name]=Angle(self.table[lat_name],unit=coord_units)
+            self.lat_name=lat_name
+            self.lon_name=lon_name
+        else:
+            self.table = catalog_data
+            meta = getattr(self.table, 'meta', {})
+            self.lat_name = meta.get('LAT_NAME')
+            self.lon_name = meta.get('LON_NAME')
 
-        self.lat_name=lat_name
-        self.lon_name=lon_name
 
-    def get_api_dictionary(self ):
+    def get_api_dictionary(self, dump_string=True):
         column_lists = []
         for colname in self.table.colnames:
             column_lists.append([x if str(x) != 'nan' else None for x in self.table[colname]])
                                 
-
-        return json.dumps(dict(cat_frame=self.table.meta['FRAME'], # pyright: ignore[reportOptionalSubscript]
+        cat_dict = dict(cat_frame=self.table.meta['FRAME'], # pyright: ignore[reportOptionalSubscript]
                     cat_coord_units=self.table.meta['COORD_UNIT'], # pyright: ignore[reportOptionalSubscript]
                     cat_column_list=column_lists,
                     cat_column_names=self.table.colnames,
                     cat_column_descr=self.table.dtype.descr,
                     cat_lat_name=self.lat_name,
-                    cat_lon_name=self.lon_name))
+                    cat_lon_name=self.lon_name)
+        if dump_string:
+            return json.dumps(cat_dict)
+        else:
+            return cat_dict
+
+    def encode(self):
+        return self.get_api_dictionary(dump_string=False)
+    
+    @classmethod
+    def decode(cls, encoded_obj : str | dict[str, typing.Any], name: str | None = None) -> "ApiCatalog":
+        if isinstance(encoded_obj, str):
+            obj = json.loads(encoded_obj)
+        else:
+            obj = encoded_obj
+
+        return cls(obj, name = name)
+
+    def suggest_fn_extension(self) -> str:
+        return 'ecsv'
+
+    def write_file(self, file_path, overwrite=True):
+            # determine format from the file extension
+            self.table.write(file_path, overwrite=overwrite)
+
+    @classmethod
+    def from_file(cls, file_path: str, name: str | None = None, delimiter: str | None = None, format: str | None = None):
+        allowed_formats=['ascii','ascii.ecsv','fits']
+        if format in allowed_formats:
+            kw = {'delimiter': delimiter} if format != 'fits' and delimiter is not None else {}
+            table = Table.read(file_path, format=format, **kw)
+        else:
+            raise RuntimeError(f'Catalog file format not understood, allowed: {allowed_formats}')
+
+        return cls(table, name=name)
 
 class GWEventContours:
     def __init__(self, event_contour_dict, name='') -> None:
diff --git a/tests/test_data_products.py b/tests/test_data_products.py
@@ -18,12 +18,15 @@
                                    ODAAstropyTable, 
                                    PictureProduct,
                                    BinaryProduct,
-                                   TextLikeProduct)
+                                   TextLikeProduct,
+                                   ApiCatalog)
 from astropy import time as atime
 from astropy import units as u
 from astropy.table import Table
 from matplotlib import pyplot as plt
 
+from astropy.utils.diff import report_diff_values
+
 import base64
 import pickle
 
@@ -320,6 +323,47 @@ def test_lightcurve_data_product_write_file_roundtrip(tmp_path):
     assert np.array_equal(loaded.data_unit[1].data['FLUX'], np.array(values))
     assert np.array_equal(loaded.data_unit[1].data['ERROR'], np.array(errors))
 
+def test_catalog_data_product_write_file_roundtrip(tmp_path):
+    catalog_ecsv = """\
+# %ECSV 1.0
+# ---
+# datatype:
+# - {name: meta_ID, datatype: int64}
+# - {name: src_names, datatype: string}
+# - {name: significance, datatype: float32}
+# - {name: ra, datatype: float32}
+# - {name: dec, datatype: float32}
+# - {name: NEW_SOURCE, datatype: uint16}
+# - {name: ISGRI_FLAG, datatype: int64}
+# - {name: FLAG, datatype: int64}
+# - {name: ERR_RAD, datatype: float64}
+# meta: !!omap
+# - {FRAME: fk5}
+# - {LAT_NAME: dec}
+# - {COORD_UNIT: deg}
+# - {LON_NAME: ra}
+# schema: astropy-2.0
+meta_ID src_names significance ra dec NEW_SOURCE ISGRI_FLAG FLAG ERR_RAD
+0 "IGR J15311-3737" 0.0 0.0 0.0 0 1 0 0.0750000029802
+1 "IGR J15409-4057" 0.804126 235.259 -40.9678 0 1 0 0.00860999990255
+"""
+    ecsv_file_in = tmp_path / 'catalog.ecsv'
+    with open(ecsv_file_in, 'w') as fd:
+        fd.write(catalog_ecsv)
+
+    ecsv_file_out = tmp_path / 'catalog.ecsv'
+    fits_file_out = tmp_path / 'catalog.fits'
+
+    catalog = ApiCatalog.from_file(ecsv_file_in, name = 'example catalog', format = 'ascii.ecsv')
+    catalog.write_file(ecsv_file_out)
+    catalog.write_file(fits_file_out)
+
+    with open(ecsv_file_out, 'r') as fd:
+        assert fd.read() == catalog_ecsv
+    
+    from_fits = ApiCatalog.from_file(fits_file_out, name = 'example catalog 1', format = 'fits')
+    assert report_diff_values(from_fits.table, from_fits.table)
+
 
 def test_save_all_data_mixed_collection(tmp_path):
     table = Table({'a': [1, 2], 'b': [3, 4]})
@@ -333,7 +377,18 @@ def test_save_all_data_mixed_collection(tmp_path):
     picture = PictureProduct.from_file(str(plot_fn), name='pic')
     text = TextLikeProduct('hello', name='text')
     lc = LightCurveDataProduct.from_arrays(['2022-02-20T13:45:34', '2022-02-20T14:45:34'], fluxes=[2, 3], errors=[0.1, 0.2])
-    dc = DataCollection([astable, bin_prod, picture, text, lc])
+    ascii_catalog = """\
+meta_ID src_names significance ra dec NEW_SOURCE ISGRI_FLAG FLAG ERR_RAD
+0 "IGR J15311-3737" 0.0 0.0 0.0 0 1 0 0.0750000029802
+1 "IGR J15409-4057" 0.804126 235.259 -40.9678 0 1 0 0.00860999990255
+"""
+    catalog_fn_in = tmp_path / 'catalog.csv'
+    with open(catalog_fn_in, 'w') as fd:
+        fd.write(ascii_catalog)
+    
+    catalog = ApiCatalog.from_file(catalog_fn_in, name='catalog', format='ascii', delimiter=' ')
+
+    dc = DataCollection([astable, bin_prod, picture, text, lc, catalog])
 
     out_dir = tmp_path / 'saved'
     os.makedirs(out_dir, exist_ok=True)
@@ -345,6 +400,7 @@ def test_save_all_data_mixed_collection(tmp_path):
         out_dir / 'mixed_pic_2.png',
         out_dir / 'mixed_text_3.txt',
         out_dir / 'mixed_prod_4.fits',
+        out_dir / 'mixed_catalog_5.ecsv',
     ]
 
     for expected in expected_files:
