Coverage for tbkit/save.py: 100%

1from __future__ import annotations 

2 

3import os 

4from matplotlib.figure import Figure 

5from matplotlib.animation import FuncAnimation 

6import tbkit.error_handling as error_handling 

7 

8 

9################################# 

10# CLASS SAVE 

11################################# 

12 

13 

14class Save(): 

15 ''' 

16 Create folder and save figures / animations obtained via 

17 **plot** or **propagation**. 

18 

19 Parameters 

20 ---------- 

21 dir_name : String 

22 Name of the sub-directory the figures are stored in. 

23 dir_main : String, default None 

24 Path of the main directory. 

25 If None, figures are stored under ``'figs/'``. 

26 params : Dictionary, default None 

27 Parameters appended to file names. 

28 file_format 

29 Default value 'png'. Figure format. 

30 ''' 

31 def __init__( 

32 self, 

33 dir_name: str, 

34 dir_main: str | None = None, 

35 params: dict | None = None, 

36 file_format: str = 'png', 

37 ) -> None: 

38 error_handling.string(dir_name, 'dir_name') 

39 error_handling.string(dir_main, 'dir_main') 

40 error_handling.file_format(file_format) 

41 self.params = {} if params is None else params 

42 self.file_format = file_format 

43 if dir_main is None: 

44 self.dir_main = 'figs' 

45 else: 

46 self.dir_main = dir_main 

47 self.dir_name = os.path.join(self.dir_main, dir_name) 

48 self.create_dir() 

49 

50 def create_dir(self) -> None: 

51 ''' 

52 Create the directory the figures are stored in, if it does not exist. 

53 ''' 

54 if not os.path.exists(self.dir_main): 

55 os.makedirs(self.dir_main) 

56 if not os.path.exists(self.dir_name): 

57 os.makedirs(self.dir_name) 

58 

59 def file_name(self) -> str: 

60 ''' 

61 Create the file name. 

62 

63 Returns 

64 ------- 

65 file_name 

66 File name. 

67 ''' 

68 file_name = '' 

69 for key, val in self.params.items(): 

70 file_name += '_' + key + str(complex(val+0)).replace('.', ',') 

71 return file_name 

72 

73 def fig(self, fig: Figure, name: str) -> None: 

74 ''' 

75 Save the figure in the directory *dir_name*. 

76 

77 Parameters 

78 ---------- 

79 fig : Matplotlib fig 

80 name : String 

81 First part of the file name. 

82 ''' 

83 error_handling.fig(fig) 

84 error_handling.string(name, 'name') 

85 name_file = os.path.join(self.dir_name, 

86 name + self.file_name() + '.' + self.file_format) 

87 fig.savefig(name_file, format=self.file_format) 

88 

89 def fig_lat(self, fig: Figure, name: str) -> None: 

90 ''' 

91 Save the figure in the directory *dir_name*. 

92 

93 Parameters 

94 ---------- 

95 fig : Matplotlib fig 

96 name : String 

97 First part of the file name. 

98 ''' 

99 error_handling.fig(fig) 

100 error_handling.string(name, 'name') 

101 name_file = os.path.join(self.dir_name, name + '.' + self.file_format) 

102 fig.savefig(name_file, format=self.file_format) 

103 

104 def ani(self, ani: FuncAnimation, name: str, fps: int = 10) -> None: 

105 ''' 

106 Save the animation, as MP4, in the directory *dir_name*. 

107 

108 Parameters 

109 ---------- 

110 ani : Matplotlib FuncAnimation 

111 name : String 

112 First part of the file name. 

113 fps : Positive integer, default 10 

114 Frames per second. 

115 ''' 

116 error_handling.ani(ani) 

117 error_handling.string(name, 'name') 

118 error_handling.positive_int(fps, 'fps') 

119 name_file = os.path.join(self.dir_name, name + '.mp4') 

120 ani.save(name_file, fps=fps, extra_args=['-vcodec', 'libx264']) 

121 

122 

123# Backward-compatible lowercase alias (pre-0.2 API). 

124save = Save