Coverage for /builds/debichem-team/python-ase/ase/utils/plotting.py: 91.67%

1from typing import Optional 

2 

3import matplotlib.pyplot as plt 

4from matplotlib.axes import Axes 

5from matplotlib.figure import Figure 

6 

7 

8class SimplePlottingAxes: 

9 def __init__(self, 

10 ax: Optional[Axes] = None, 

11 show: bool = False, 

12 filename: str = None) -> None: 

13 self.ax = ax 

14 self.show = show 

15 self.filename = filename 

16 self.figure: Optional[Figure] = None 

17 

18 def __enter__(self) -> Axes: 

19 if self.ax is None: 

20 self.figure, self.ax = plt.subplots() 

21 else: 

22 self.figure = self.ax.get_figure() 

23 

24 return self.ax 

25 

26 def __exit__(self, exc_type, exc_val, exc_tb) -> None: 

27 if exc_type is None: 

28 # If there was no exception, display/write the plot as appropriate 

29 if self.figure is None: 

30 raise Exception("Something went wrong initializing matplotlib " 

31 "figure") 

32 if self.show: 

33 self.figure.show() 

34 if self.filename is not None: 

35 self.figure.savefig(self.filename) 

36 

37 return