Skip to content

Commit

Permalink
extend declarative colorbar functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
kgoebber committed Sep 6, 2023
1 parent c0826bc commit ce1c28d
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 8 deletions.
22 changes: 14 additions & 8 deletions src/metpy/plots/declarative.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from traitlets import (Any, Bool, Float, HasTraits, Instance, Int, List, observe, TraitError,
Tuple, Unicode, Union, validate)
from traitlets import (Any, Bool, Dict, Float, HasTraits, Instance, Int, List, observe,
TraitError, Tuple, Unicode, Union, validate)

from . import ctables, wx_symbols
from ._mpl import TextCollection
Expand Down Expand Up @@ -825,8 +825,12 @@ def draw(self):
if getattr(self, 'handle', None) is None:
self._build()
if getattr(self, 'colorbar', None) is not None:
cbar = self.parent.ax.figure.colorbar(
self.handle, orientation=self.colorbar, pad=0, aspect=50)
if isinstance(self.colorbar, dict):
cbar = self.parent.ax.figure.colorbar(
self.handle, **self.colorbar)
else:
cbar = self.parent.ax.figure.colorbar(
self.handle, orientation=self.colorbar, pad=0, aspect=50)
cbar.ax.tick_params(labelsize=self.colorbar_fontsize)
self._need_redraw = False

Expand Down Expand Up @@ -878,13 +882,15 @@ class ColorfillTraits(MetPyHasTraits):
`matplotlib.colors.Normalize` instance for plotting.
"""

colorbar = Unicode(default_value=None, allow_none=True)
colorbar = Union([Unicode(default_value=None, allow_none=True), Dict()])
colorbar.__doc__ = """A string (horizontal/vertical) on whether to add a colorbar to the
plot.
To add a colorbar associated with the plot, set the trait to ``horizontal`` or
``vertical``,specifying the orientation of the produced colorbar. The default value is
``None``.
To add a colorbar associated with the plot, you can either set the trait with a string of
``horizontal`` or ``vertical``, which specifies the orientation of the produced colorbar
and uses pre-defined defaults for aspect and pad. Alternatively, you can set a dictionary
of keyword argument values valid for a Matplotlib colorbar to specify how the colorbar will
be plotted. The default value is ``None``.
"""

colorbar_fontsize = Union([Int(), Float(), Unicode()], allow_none=True, default_value=None)
Expand Down
Binary file added tests/plots/baseline/test_colorbar_kwargs.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
25 changes: 25 additions & 0 deletions tests/plots/test_declarative.py
Original file line number Diff line number Diff line change
Expand Up @@ -693,6 +693,31 @@ def test_colorfill_horiz_colorbar(cfeature):
return pc.figure


@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.02)
def test_colorbar_kwargs(cfeature):
"""Test that we can use ContourFillPlot with specifying colorbar kwargs."""
data = xr.open_dataset(get_test_data('narr_example.nc', as_file_obj=False))

contour = FilledContourPlot()
contour.data = data
contour.level = 700 * units.hPa
contour.field = 'Temperature'
contour.colormap = 'coolwarm'
contour.colorbar = {'orientation': 'horizontal', 'aspect': 60, 'pad': 0.05}

panel = MapPanel()
panel.area = (-110, -60, 25, 55)
panel.layers = []
panel.plots = [contour]

pc = PanelContainer()
pc.panel = panel
pc.size = (8, 8)
pc.draw()

return pc.figure


@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.02)
def test_colorfill_no_colorbar(cfeature):
"""Test that we can use ContourFillPlot with no colorbar."""
Expand Down

0 comments on commit ce1c28d

Please sign in to comment.