1from typing import Optional 

2 

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

4if False: 

5 import matplotlib.axes 

6 import matplotlib.figure 

7 

8 

9class SimplePlottingAxes: 

10 def __init__(self, 

11 ax: 'matplotlib.axes.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 = None # type: Optional[matplotlib.figure.Figure] 

18 

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

20 if self.ax is None: 

21 import matplotlib.pyplot as plt 

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

23 else: 

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

25 

26 return self.ax 

27 

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

29 if exc_type is None: 

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

31 if self.figure is None: 

32 raise Exception("Something went wrong initializing matplotlib " 

33 "figure") 

34 if self.show: 

35 self.figure.show() 

36 if self.filename is not None: 

37 self.figure.savefig(self.filename) 

38 

39 return None