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 = self.dir_main + dir_name 

42 self.create_dir() 

43 

44 def create_dir(self) -> None: 

45 ''' 

46 Create the directory to store the figures exists. 

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 defined by the method *dir_name()*. 

68 

69 :param fig: Matplotlib fig. 

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

71 ''' 

72 error_handling.fig(fig) 

73 error_handling.string(name, 'name') 

74 name_file = self.dir_name + '/' + name + self.file_name() + '.' + self.file_format 

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

76 

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

78 ''' 

79 Save the figure in the directory defined by the method *dir_name()*. 

80 

81 :param fig: Matplotlib fig. 

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

83 ''' 

84 error_handling.fig(fig) 

85 error_handling.string(name, 'name') 

86 name_file = self.dir_name + '/' + name + '.' + self.file_format 

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

88 

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

90 error_handling.ani(ani) 

91 error_handling.string(name, 'name') 

92 error_handling.positive_int(fps, 'fps') 

93 name_file = self.dir_name + '/' + name + '.mp4' 

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

95 

96 

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

98save = Save