Coverage for /builds/ericyuan00000/ase/ase/utils/plotting.py: 91.30%

1# fmt: off 

2 

3from typing import Any, Optional 

4 

5import matplotlib.pyplot as plt 

6from matplotlib.axes import Axes 

7 

8 

9class SimplePlottingAxes: 

10 def __init__(self, 

11 ax: Optional[Axes] = None, 

12 show: bool = False, 

13 filename: str = None) -> None: 

14 self.ax = ax 

15 self.show = show 

16 self.filename = filename 

17 self.figure: Any = None # Don't know about Figure/SubFigure etc 

18 

19 def __enter__(self) -> Axes: 

20 if self.ax is None: 

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

22 else: 

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

24 

25 return self.ax 

26 

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

28 if exc_type is None: 

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

30 if self.figure is None: 

31 raise Exception("Something went wrong initializing matplotlib " 

32 "figure") 

33 if self.show: 

34 self.figure.show() 

35 if self.filename is not None: 

36 self.figure.savefig(self.filename) 

37 

38 return