From 382fe968375dce03f9c7b02524c9b6c1bb1381df Mon Sep 17 00:00:00 2001 From: Thomas A Caswell Date: Sun, 27 Mar 2016 17:52:19 -0400 Subject: [PATCH] MNT: use stdlib tools in allow_rasterization Use a context manager to mange enabling/disabling the rasterization. This will now try to un-wrap the rasterization context state if an exception occurs. Use functools.wraps instead of doing it 'manually'. --- lib/matplotlib/artist.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py index 37c78e8a859..0624c132936 100644 --- a/lib/matplotlib/artist.py +++ b/lib/matplotlib/artist.py @@ -15,7 +15,8 @@ from .transforms import (Bbox, IdentityTransform, TransformedBbox, TransformedPatchPath, TransformedPath, Transform) from .path import Path - +from functools import wraps +from contextlib import contextmanager # Note, matplotlib artists use the doc strings for set and get # methods to enable the introspection methods of setp and getp. Every # set_* method should have a docstring containing the line @@ -42,31 +43,30 @@ def allow_rasterization(draw): other setup function calls, such as starting and flushing a mixed-mode renderer. """ - def before(artist, renderer): + @contextmanager + def with_rasterized(artist, renderer): + if artist.get_rasterized(): renderer.start_rasterizing() if artist.get_agg_filter() is not None: renderer.start_filter() - def after(artist, renderer): - - if artist.get_agg_filter() is not None: - renderer.stop_filter(artist.get_agg_filter()) + try: + yield + finally: + if artist.get_agg_filter() is not None: + renderer.stop_filter(artist.get_agg_filter()) - if artist.get_rasterized(): - renderer.stop_rasterizing() + if artist.get_rasterized(): + renderer.stop_rasterizing() # the axes class has a second argument inframe for its draw method. + @wraps(draw) def draw_wrapper(artist, renderer, *args, **kwargs): - before(artist, renderer) - draw(artist, renderer, *args, **kwargs) - after(artist, renderer) - - # "safe wrapping" to exactly replicate anything we haven't overridden above - draw_wrapper.__name__ = draw.__name__ - draw_wrapper.__dict__ = draw.__dict__ - draw_wrapper.__doc__ = draw.__doc__ + with with_rasterized(artist, renderer): + return draw(artist, renderer, *args, **kwargs) + draw_wrapper._supports_rasterization = True return draw_wrapper