save_all_data() and related tests

Author: dsavchenko

oda_api/api.py

@@ -35,7 +35,7 @@
 from . import __version__
 from . import colors as C
 from . import custom_formatters
-from .data_products import (ApiCatalog, BinaryData, BinaryProduct,
+from .data_products import (ApiCatalog, BinaryData, BinaryProduct, DataProduct,
                             GWContoursDataProduct, NumpyDataProduct,
                             ODAAstropyTable, PictureProduct, TextLikeProduct)
 
@@ -1196,11 +1196,17 @@ def __repr__(self):
 
 class DataCollection(object):
 
-    def __init__(self, data_list, add_meta_to_name=None, instrument=None, product=None, request_job_id=None):
+    def __init__(self, 
+                 data_list: list[DataProduct], 
+                 add_meta_to_name: list[str] | None = None, 
+                 instrument: str | None = None, 
+                 product: str | None =None, 
+                 request_job_id: str | None=None
+                 ):
         if add_meta_to_name is None:
             add_meta_to_name = ['src_name', 'product']
-        self._p_list = []
-        self._n_list = []
+        self._p_list: list[DataProduct] = []
+        self._n_list: list[str] = []
         self.request_job_id = request_job_id
         for ID, data in enumerate(data_list):
 
@@ -1260,15 +1266,21 @@ def _build_prod_name(self, prod, name, add_meta_to_name):
                         name += '_' + s.strip()
         return name, oda_api.misc_helpers.clean_var_name(name)
 
-    def save_all_data(self, prenpend_name=None):
+    def save_all_data(self, prenpend_name = None, overwrite=True):
+        # NOTE: prepend_name also determines file path
         for pname, prod in zip(self._n_list, self._p_list, strict=False):
+            if not isinstance(prod, DataProduct):
+                logger.warning(f"Writing on disk is not implemented for product {pname} of type {pname.__class__.__name__}, skipping.")
+                continue
             if prenpend_name is not None:
                 file_name = prenpend_name + '_' + pname
             else:
                 file_name = pname
 
-            file_name = file_name + '.fits'
-            prod.write_fits_file(file_name)
+            fn_extension = prod.suggest_fn_extension()
+            file_name = f"{file_name}.{fn_extension}"
+
+            prod.write_file(file_name, overwrite=overwrite)
 
     def save(self, file_name):
         pickle.dump(self, open(file_name, 'wb'),
@@ -1363,7 +1375,7 @@ def from_response_json(cls, res_json, instrument, product):
         d = cls(data, instrument=instrument, product=product, request_job_id=request_job_id)
         for p in d._p_list:
             if hasattr(p, 'meta_data') is False and hasattr(p, 'meta') is True:
-                p.meta_data = p.meta
+                p.meta_data = p.meta  # type:ignore
 
         return d
 

oda_api/data_products.py

@@ -49,6 +49,9 @@ def _chekc_enc_data(data):
 
 
 class DataProduct(ABC):
+    name: str | None
+    meta_data: dict
+
     @abstractmethod
     def encode(self, *args, **kwargs) -> str | dict[str, typing.Any]: ...
 
@@ -192,7 +195,7 @@ def __init__(self, bin_data: bytes, name: str | None = None):
     def suggest_fn_extension(self) -> str:
         suggested_extension = puremagic.from_string(self.bin_data)
         if suggested_extension:
-            return suggested_extension
+            return suggested_extension.strip('.')
         return 'bin'
 
     def encode(self):

tests/test_data_products.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+import logging
 
 import pytest
 import json
@@ -8,6 +9,7 @@
 import time
 import os
 import typing
+from oda_api.api import DataCollection
 from oda_api.json import CustomJSONEncoder
 import filecmp
 
@@ -317,3 +319,51 @@ def test_lightcurve_data_product_write_file_roundtrip(tmp_path):
     loaded = NumpyDataProduct.from_file(str(out_fn))
     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_save_all_data_mixed_collection(tmp_path):
+    table = Table({'a': [1, 2], 'b': [3, 4]})
+    astable = ODAAstropyTable(table, name='table')
+    bin_prod = BinaryProduct.from_file('tests/test_data/lc.fits', name='binprd')
+
+    plot_fn = tmp_path / 'plot.png'
+    plt.plot([1, 2], [1, 0])
+    plt.savefig(str(plot_fn))
+    plt.close()
+    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])
+
+    out_dir = tmp_path / 'saved'
+    os.makedirs(out_dir, exist_ok=True)
+    dc.save_all_data(prenpend_name=str(out_dir / 'mixed'))
+
+    expected_files = [
+        out_dir / 'mixed_table_0.fits',
+        out_dir / 'mixed_binprd_1.fits',
+        out_dir / 'mixed_pic_2.png',
+        out_dir / 'mixed_text_3.txt',
+        out_dir / 'mixed_prod_4.fits',
+    ]
+
+    for expected in expected_files:
+        assert expected.exists()
+
+
+def test_save_all_data_empty_collection(tmp_path):
+    dc = DataCollection([])
+    dc.save_all_data(prenpend_name=str(tmp_path / 'empty'))
+    assert not any(tmp_path.iterdir())
+
+
+def test_save_all_data_skips_unsupported_product(caplog, tmp_path):
+    caplog.set_level(logging.WARNING)
+    text = TextLikeProduct('hello', name='text')
+    unsupported = object()
+    dc = DataCollection([text, unsupported])
+
+    dc.save_all_data(prenpend_name=str(tmp_path / 'skipped'))
+
+    assert 'Writing on disk is not implemented for' in caplog.text
+    assert (tmp_path / 'skipped_text_0.txt').exists()