from __future__ import annotations
import typing
from copy import deepcopy
from warnings import warn
import numpy as np
import pandas as pd
from ..exceptions import PlotnineError, PlotnineWarning
from ..utils import Registry, is_string
from .guide import guide as guide_class
if typing.TYPE_CHECKING:
from typing import Literal
from plotnine.typing import TupleFloat2
# Terminology
# -----------
# - A guide is either a legend or colorbar.
#
# - A guide definition (gdef) is an instantiated guide as it
# is used in the process of creating the legend
#
# - The guides class holds all guides that will appear in the
# plot
#
# - A guide box is a fully drawn out guide.
# It is of subclass matplotlib.offsetbox.Offsetbox
[docs]class guides(dict):
"""
Guides for each scale
Used to assign a particular guide to an aesthetic(s).
Parameters
----------
kwargs : dict
aesthetic - guide pairings. e.g
``color=guide_colorbar()``
"""
# Determined from the theme when the guides are
# getting built
position: Literal["left", "right", "top", "bottom"] | TupleFloat2
box_direction: Literal["horizontal", "vertical", "auto"]
box_align: Literal["left", "right", "top", "bottom", "center", "auto"]
box_margin: int
spacing: float
def __init__(self, **kwargs):
aes_names = {
"alpha",
"color",
"fill",
"linetype",
"shape",
"size",
"stroke",
}
if "colour" in kwargs:
kwargs["color"] = kwargs.pop("colour")
dict.__init__(
self, ((ae, kwargs[ae]) for ae in kwargs if ae in aes_names)
)
def __radd__(self, gg):
"""
Add guides to the plot
Parameters
----------
gg : ggplot
ggplot object being created
Returns
-------
out : gglot
ggplot object with guides. If *inplace*
is **False** this is a copy of the original
ggplot object.
"""
new_guides = {}
for k in self:
new_guides[k] = deepcopy(self[k])
gg.guides.update(new_guides)
return gg
def build(self, plot):
"""
Build the guides
Parameters
----------
plot : ggplot
ggplot object being drawn
Returns
-------
box : matplotlib.offsetbox.Offsetbox | None
A box that contains all the guides for the plot.
If there are no guides, **None** is returned.
"""
_property = plot.theme.themeables.property
self.box_direction = _property("legend_box")
self.position = _property("legend_position")
self.box_align = _property("legend_box_just")
self.box_margin = _property("legend_box_margin")
self.spacing = _property("legend_spacing")
if self.position == "none":
return # No Legend
# Direction
if self.box_direction == "auto":
if self.position in ("right", "left"):
self.box_direction = "vertical"
else:
self.box_direction = "horizontal"
# Justification of legend boxes
if self.box_align == "auto":
if self.position in ("right", "left"):
self.box_align = "left"
else:
self.box_align = "right"
gdefs = self.train(plot)
if not gdefs:
return
gdefs = self.merge(gdefs)
gdefs = self.create_geoms(gdefs, plot)
if not gdefs:
return
gboxes = self.draw(gdefs, plot.theme)
bigbox = self.assemble(gboxes, gdefs, plot.theme)
return bigbox
def train(self, plot):
"""
Compute all the required guides
Parameters
----------
plot : ggplot
ggplot object
Returns
-------
gdefs : list
Guides for the plots
"""
gdefs = []
for scale in plot.scales:
for output in scale.aesthetics:
# The guide for aesthetic 'xxx' is stored
# in plot.guides['xxx']. The priority for
# the guides depends on how they are created
# 1. ... + guides(xxx=guide_blah())
# 2. ... + scale_xxx(guide=guide_blah())
# 3. default(either guide_legend or guide_colorbar
# depending on the scale type)
# output = scale.aesthetics[0]
guide = self.get(output, scale.guide)
if guide is None or guide is False:
continue
# check the validity of guide.
# if guide is character, then find the guide object
guide = self.validate(guide)
# check the consistency of the guide and scale.
if (
"any" not in guide.available_aes
and scale.aesthetics[0] not in guide.available_aes
):
raise PlotnineError(
f"{guide.__class__.__name__} cannot be used for "
f"{scale.aesthetics}"
)
# title
if not hasattr(guide, "title"):
if scale.name:
guide.title = scale.name
else:
guide.title = getattr(plot.labels, output)
if guide.title is None:
warn(
f"Cannot generate legend for the {output!r} "
"aesthetic. Make sure you have mapped a "
"variable to it",
PlotnineWarning,
)
# each guide object trains scale within the object,
# so Guides (i.e., the container of guides)
# need not to know about them
guide = guide.train(scale, output)
if guide is not None:
gdefs.append(guide)
return gdefs
def validate(self, guide):
"""
Validate guide object
"""
if is_string(guide):
guide = Registry[f"guide_{guide}"]()
if not isinstance(guide, guide_class):
raise PlotnineError(f"Unknown guide: {guide}")
return guide
def merge(self, gdefs):
"""
Merge overlapped guides
For example::
from plotnine import *
gg = ggplot(mtcars, aes(y='wt', x='mpg', colour='factor(cyl)'))
gg = gg + stat_smooth(aes(fill='factor(cyl)'), method='lm')
gg = gg + geom_point()
gg
This would create two guides with the same hash
"""
# group guide definitions by hash, and
# reduce each group to a single guide
# using the guide.merge method
definitions = pd.DataFrame(
{"gdef": gdefs, "hash": [g.hash for g in gdefs]}
)
grouped = definitions.groupby("hash", sort=False)
gdefs = []
for name, group in grouped:
# merge
gdef = group["gdef"].iloc[0]
for g in group["gdef"].iloc[1:]:
gdef = gdef.merge(g)
gdefs.append(gdef)
return gdefs
def create_geoms(self, gdefs, plot):
"""
Add geoms to the guide definitions
"""
new_gdefs = []
for gdef in gdefs:
gdef = gdef.create_geoms(plot)
if gdef:
new_gdefs.append(gdef)
return new_gdefs
def draw(self, gdefs, theme):
"""
Draw out each guide definition
Parameters
----------
gdefs : list of guide_legend|guide_colorbar
guide definitions
theme : theme
Plot theme
Returns
-------
out : list of matplotlib.offsetbox.Offsetbox
A drawing of each legend
"""
for g in gdefs:
g._set_defaults(theme)
return [g.draw() for g in gdefs]
def assemble(self, gboxes, gdefs, theme):
"""
Put together all the guide boxes
Parameters
----------
gboxes : list
List of :class:`~matplotlib.offsetbox.Offsetbox`,
where each item is a legend for a single aesthetic.
gdefs : list of guide_legend|guide_colorbar
guide definitions
theme : theme
Plot theme
Returns
-------
box : OffsetBox
A box than can be placed onto a plot
"""
from matplotlib.offsetbox import HPacker, VPacker
# place the guides according to the guide.order
# 0 do not sort
# 1-99 sort
for gdef in gdefs:
if gdef.order == 0:
gdef.order = 100
elif not 0 <= gdef.order <= 99:
raise PlotnineError(
"'order' for a guide should be " "between 0 and 99"
)
orders = [gdef.order for gdef in gdefs]
idx = np.argsort(orders)
gboxes = [gboxes[i] for i in idx]
# direction when more than legend
if self.box_direction == "vertical":
packer = VPacker
elif self.box_direction == "horizontal":
packer = HPacker
else:
raise PlotnineError(
"'legend_box' should be either " "'vertical' or 'horizontal'"
)
box = packer(
children=gboxes,
align=self.box_align,
pad=self.box_margin,
sep=self.spacing,
)
return box