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 :param dir_name: String. Name of the sub-directory the figures are stored in. 

20 :param dir_main: String. Default value None. Path of the main directory. 

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

22 :param params: Dictionary. Default value None. Parameters appended to file names. 

23 :param file_format: Default value 'png'. Figure format. 

24 ''' 

25 def __init__( 

26 self, 

27 dir_name: str, 

28 dir_main: str | None = None, 

29 params: dict | None = None, 

30 file_format: str = 'png', 

31 ) -> None: 

32 error_handling.string(dir_name, 'dir_name') 

33 error_handling.string(dir_main, 'dir_main') 

34 error_handling.file_format(file_format) 

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

36 self.file_format = file_format 

37 if dir_main is None: 

38 self.dir_main = 'figs' 

39 else: 

40 self.dir_main = dir_main 

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

42 self.create_dir() 

43 

44 def create_dir(self) -> None: 

45 ''' 

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

47 ''' 

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

49 os.makedirs(self.dir_main) 

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

51 os.makedirs(self.dir_name) 

52 

53 def file_name(self) -> str: 

54 ''' 

55 Create the file name. 

56 

57 :returns: 

58 * **file_name** -- File name. 

59 ''' 

60 file_name = '' 

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

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

63 return file_name 

64 

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

66 ''' 

67 Save the figure in the directory *dir_name*. 

68 

69 :param fig: Matplotlib fig. 

70 :param name: String. First part of the file name. 

71 ''' 

72 error_handling.fig(fig) 

73 error_handling.string(name, 'name') 

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

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

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

77 

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

79 ''' 

80 Save the figure in the directory *dir_name*. 

81 

82 :param fig: Matplotlib fig. 

83 :param name: String. First part of the file name. 

84 ''' 

85 error_handling.fig(fig) 

86 error_handling.string(name, 'name') 

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

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

89 

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

91 ''' 

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

93 

94 :param ani: Matplotlib FuncAnimation. 

95 :param name: String. First part of the file name. 

96 :param fps: Positive integer. Default value 10. Frames per second. 

97 ''' 

98 error_handling.ani(ani) 

99 error_handling.string(name, 'name') 

100 error_handling.positive_int(fps, 'fps') 

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

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

103 

104 

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

106save = Save