Coverage for /home/runner/work/nr-catalog-tools/nr-catalog-tools/nrcatalogtools/catalog.py: 69%

1import os 

2from abc import ABC, abstractmethod 

3 

4import sxs 

5from nrcatalogtools import waveform 

6 

7 

8class CatalogABC(ABC): 

9 @abstractmethod 

10 def waveform_filename_from_simname(self, sim_name): 

11 raise NotImplementedError() 

12 

13 @abstractmethod 

14 def waveform_filepath_from_simname(self, sim_name): 

15 raise NotImplementedError() 

16 

17 @abstractmethod 

18 def metadata_filename_from_simname(self, sim_name): 

19 raise NotImplementedError() 

20 

21 @abstractmethod 

22 def metadata_filepath_from_simname(self, sim_name): 

23 raise NotImplementedError() 

24 

25 @abstractmethod 

26 def download_waveform_data(self, sim_name): 

27 raise NotImplementedError() 

28 

29 @abstractmethod 

30 def waveform_url_from_simname(self, sim_name): 

31 raise NotImplementedError() 

32 

33 

34class CatalogBase(CatalogABC, sxs.Catalog): 

35 def __init__(self, *args, **kwargs) -> None: 

36 sxs.Catalog.__init__(self, *args, **kwargs) 

37 

38 def get(self, sim_name): 

39 if sim_name not in self.simulations_dataframe.index.to_list(): 

40 raise IOError( 

41 f"Simulation {sim_name} not found in catalog." 

42 f"Please check that it exists" 

43 ) 

44 filepath = self.waveform_filepath_from_simname(sim_name) 

45 if not os.path.exists(filepath) or os.path.getsize(filepath) == 0: 

46 if self._verbosity > 1: 

47 print( 

48 f"..As data does not exist in cache:" 

49 f" (in {filepath}),\n" 

50 f"..we will now download it from" 

51 " {}".format(self.waveform_url_from_simname(sim_name)) 

52 ) 

53 self.download_waveform_data(sim_name) 

54 metadata = self.get_metadata(sim_name) 

55 if type(metadata) is not dict and hasattr(metadata, "to_dict"): 

56 metadata = metadata.to_dict() 

57 return waveform.WaveformModes.load_from_h5(filepath, metadata=metadata) 

58 

59 def get_metadata(self, sim_name): 

60 sim_dict = self.simulations 

61 if sim_name not in list(sim_dict.keys()): 

62 raise IOError( 

63 f"Simulation {sim_name} not found in catalog." 

64 f"Please check that it exists" 

65 ) 

66 return sim_dict[sim_name] 

67 

68 def set_attribute_in_waveform_data_file(self, sim_name, attr_name, attr_value): 

69 """Set attributes in the HDF5 file holding waveform data for a given 

70 simulation 

71 

72 Args: 

73 sim_name (str): Name/Tag of the simulation 

74 attr_name (str): Name of the attribute to set 

75 attr_value (any/serializable): Value of the attribute 

76 """ 

77 import h5py 

78 

79 file_path = self.waveform_filepath_from_simname(sim_name) 

80 with h5py.File(file_path, "a") as fp: 

81 if attr_name not in fp.attrs: 

82 fp.attrs[attr_name] = attr_value