diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2cb8f2d..1e6f893 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -4,6 +4,12 @@ Added ----- * `psyplot.data.open_dataset` now decodes grid_mappings attributes, see `#17 `__ +* psyplot projects now support the with syntax, e.g. something like:: + + with psy.plot.mapplot('file.nc') as sp: + sp.export('output.png') + + sp will be closed automatically (see commit `ee7415b `__) Changed ------- diff --git a/psyplot/project.py b/psyplot/project.py index 6b84134..5273e05 100755 --- a/psyplot/project.py +++ b/psyplot/project.py @@ -390,6 +390,12 @@ def __call__(self, *args, **kwargs): ret.main = self.main return ret + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close(True, True, True) + @_first_main def extend(self, *args, **kwargs): len0 = len(self) diff --git a/tests/test_project.py b/tests/test_project.py index fa02c32..0d64e60 100644 --- a/tests/test_project.py +++ b/tests/test_project.py @@ -54,6 +54,19 @@ def tearDown(self): pass self._created_files.clear() + def test_with(self): + """Test __enter__ and __exit__ methods""" + psy.register_plotter('test_plotter', import_plotter=True, + module='test_plotter', plotter_name='TestPlotter') + self.assertFalse(psy.gcp(True)) + self.assertFalse(psy.gcp()) + with psy.plot.test_plotter(bt.get_file('test-t2m-u-v.nc')) as sp: + self.assertTrue(psy.gcp(True)) + self.assertTrue(psy.gcp()) + + self.assertFalse(psy.gcp(True)) + self.assertFalse(psy.gcp()) + def test_save_and_load_01_simple(self): """Test the saving and loading of a Project""" psy.register_plotter('test_plotter', import_plotter=True,