Coverage for src/susi/io/multi_hdu_writer.py: 100%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4Module provides `MultiHDUWriter` 

5 

6@author: hoelken 

7""" 

8 

9from astropy.io import fits 

10import numpy as np 

11 

12from .fits_header import Card 

13 

14 

15class MultiHDUWriter: 

16 """ 

17 ## MultiHDUWriter 

18 Writes the given datasets / header(s) to individual HDUs in a FITS file 

19 """ 

20 

21 def __init__(self, path): 

22 #: path of the fits to read/write 

23 self.path = path 

24 #: list of the data to set. 

25 self.data = [] 

26 #: A list of header information per dataset (align indices with data) 

27 self.header = [] 

28 #: The header to apply to all datasets 

29 self.global_header = [] 

30 #: Flag to tell if it has been saved 

31 self.saved = False 

32 

33 def custom_header(self, card: Card, idx: int = None): 

34 """ 

35 Appends the given card to the current header 

36 

37 ### Params 

38 - card: the card to append 

39 - idx:Optional, defines the data-id to add the header to. 

40 If not set (default) the card is added to the global header for all HDUs 

41 

42 ### Returns 

43 `self` 

44 """ 

45 if idx is None: 

46 self.global_header.append(card) 

47 else: 

48 self.header[idx].append(card) 

49 self.saved = False 

50 return self 

51 

52 def append(self, data: np.array, header: list = None): 

53 """ 

54 Append data to the 

55 

56 ### Params 

57 - data: the data array to append to the HDUList 

58 - header: Optional list of custom header `Cards`. 

59 """ 

60 self.data.append(data) 

61 self.header.append([] if header is None else header) 

62 self.saved = False 

63 

64 def write_to_disk(self, overwrite: bool = False, astype: str = 'float32'): 

65 """ 

66 Writes the data to the configured path 

67 

68 ### Params 

69 - overwrite: bool, optional 

70 If `True`, overwrite the output file if it exists. Raises an 

71 `OSError` if `False` and the output file exists. (Default=`False`) 

72 - astype: string, optional 

73 Define the datatype of the HDU data array. Default: float32 

74 

75 ### Returns 

76 The HDUList written to file 

77 """ 

78 hdul = fits.HDUList(self.__gen_hdu(0, astype)) 

79 for idx in range(1, len(self.data)): 

80 hdul.append(self.__gen_hdu(idx, astype)) 

81 hdul.writeto(self.path, overwrite=overwrite) 

82 self.saved = True 

83 return hdul 

84 

85 def __gen_hdu(self, idx, astype): 

86 hdu = fits.PrimaryHDU(self.data[idx].astype(astype)) 

87 for card in self.global_header: 

88 hdu.header.append(card.to_card()) 

89 for card in self.header[idx]: 

90 hdu.header.append(card.to_card()) 

91 return hdu 

92 

93 def __repr__(self): 

94 txt = '{}: {}\n'.format(self.__class__.__name__, self.path) 

95 txt += '{:<21} = {}\n{:<21} = {}\n'.format('Saved', self.saved, 'HDUS', len(self.data)) 

96 txt += '==== Global Header ====\n' 

97 txt += '\n'.join(['{:<21} = {} | {}'.format(c.name, c.value, c.comment) for c in self.global_header]) 

98 for i in range(len(self.data)): 

99 txt += '\n==== HDU [{}] ====\n'.format(i) 

100 txt += '{:<21} = {}\n'.format('Data Shape', self.data[i].shape) 

101 txt += '\n'.join(['{:<21} = {} | {}'.format(c.name, c.value, c.comment) for c in self.header[i]]) 

102 return txt