Coverage for /builds/kinetik161/ase/ase/utils/plotting.py: 90.91%

1# This import is for the benefit of type-checking / mypy 

2if False: 

3 import matplotlib.axes 

4 import matplotlib.figure 

5 

6 

7class SimplePlottingAxes: 

8 def __init__(self, 

9 ax: 'matplotlib.axes.Axes' = None, 

10 show: bool = False, 

11 filename: str = None) -> None: 

12 self.ax = ax 

13 self.show = show 

14 self.filename = filename 

15 self.figure = None 

16 

17 def __enter__(self) -> 'matplotlib.axes.Axes': 

18 if self.ax is None: 

19 import matplotlib.pyplot as plt 

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 None