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

1from typing import Any, Optional 

2 

3import matplotlib.pyplot as plt 

4from matplotlib.axes import Axes 

5 

6 

7class SimplePlottingAxes: 

8 def __init__( 

9 self, 

10 ax: Optional[Axes] = None, 

11 show: bool = False, 

12 filename: str = None, 

13 ) -> 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( 

32 'Something went wrong initializing matplotlib figure' 

33 ) 

34 if self.show: 

35 self.figure.show() 

36 if self.filename is not None: 

37 self.figure.savefig(self.filename) 

38 

39 return