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

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 #: List of headers for each hdu, if set self.header and self.global_header are ignored 

27 self.headers = None 

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

29 self.header = [] 

30 #: The header to apply to all datasets 

31 self.global_header = [] 

32 #: Flag to tell if it has been saved 

33 self.saved = False 

34 

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

36 """ 

37 Appends the given card to the current header 

38 

39 ### Params 

40 - card: the card to append 

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

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

43 

44 ### Returns 

45 `self` 

46 """ 

47 if idx is None: 

48 self.global_header.append(card) 

49 else: 

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

51 self.saved = False 

52 return self 

53 

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

55 """ 

56 Append data to the 

57 

58 ### Params 

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

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

61 """ 

62 self.data.append(data) 

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

64 self.saved = False 

65 

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

67 """ 

68 Writes the data to the configured path 

69 

70 ### Params 

71 - overwrite: bool, optional 

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

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

74 - astype: string, optional 

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

76 

77 ### Returns 

78 The HDUList written to file 

79 """ 

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

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

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

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

84 self.saved = True 

85 return hdul 

86 

87 def __gen_hdu(self, idx, astype): 

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

89 if self.headers is not None: 

90 hdu.header = self.headers[idx] 

91 else: 

92 for card in self.global_header: 

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

94 for card in self.header[idx]: 

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

96 return hdu 

97 

98 def __repr__(self): 

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

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

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

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

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

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

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

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

107 return txt