Source code for genno.compat.plotnine.plot

import logging
from abc import ABC, abstractmethod
from collections.abc import Hashable, Sequence
from pathlib import Path
from typing import Any, Optional
from warnings import warn

import plotnine as p9

import genno
from genno.compat.pandas import disable_copy_on_write
from genno.core.computer import Computer
from genno.core.key import KeyLike

log = logging.getLogger(__name__)




class Plot(ABC):
    """Class for plotting using :mod:`plotnine`."""

    #: File name base for saving the plot.
    basename = ""
    #: File extension; determines file format.
    suffix = ".pdf"
    #: Path for file output. If it is not set, :meth:`save` will populate it with a
    #: value constructed from :py:`config["output_dir"]`, :attr:`basename`, and
    #: :attr:`suffix`. The implementation of :meth:`generate` in a Plot sub-class may
    #: assign any other value, for instance one constructed at runtime from the
    #: :attr:`inputs`.
    path: Optional[Path] = None
    #: :class:`Keys <.Key>` referring to :class:`Quantities <.Quantity>` or other inputs
    #: accepted by :meth:`generate`.
    inputs: Sequence[Hashable] = []
    #: Keyword arguments for :any:`plotnine.ggplot.save`.
    save_args: dict[str, Any] = dict(verbose=False)

    # TODO add static geoms automatically in generate()
    __static: Sequence = []



    def save(self, config, *args, **kwargs) -> Optional[Path]:
        """Prepare data, call :meth:`.generate`, and save to file.

        This method is used as the callable in the task generated by :meth:`.add_tasks`.

        .. versionadded:: 1.24.1

           This method uses :func:`.disable_copy_on_write` to work around
           `has2k1/mizani#38 <https://github.com/has2k1/mizani/issues/38>`_. This may
           cause issues if other computations (for instance, of the inputs to the Plot)
           rely on Pandas' copy-on-write behaviour being enabled.
        """
        self.path = self.path or (
            config["output_dir"] / f"{self.basename}{self.suffix}"
        )

        missing = tuple(filter(lambda arg: isinstance(arg, str), args))
        if len(missing):
            log.error(
                f"Missing input(s) {missing!r} to plot {self.basename!r}; no output"
            )
            return None

        # Convert Quantity arguments to pd.DataFrame for use with plotnine
        _args = map(
            lambda arg: arg
            if not isinstance(arg, genno.Quantity)
            else arg.to_series()
            .rename(arg.name or "value")
            .reset_index()
            .assign(unit=f"{arg.units:~}"),
            args,
        )

        plot_or_plots = self.generate(*_args, **kwargs)

        if not plot_or_plots:
            log.info(
                f"{self.__class__.__name__}.generate() returned {plot_or_plots!r}; no "
                "output"
            )
            return None

        log.info(f"Save to {self.path}")

        with disable_copy_on_write(f"{__name__}.Plot.save()"):
            try:
                # Single plot
                plot_or_plots.save(self.path, **self.save_args)
            except AttributeError:
                # Iterator containing 0 or more plots
                p9.save_as_pdf_pages(plot_or_plots, self.path, **self.save_args)

        return self.path




    @classmethod
    def make_task(cls, *inputs):
        """Return a task :class:`tuple` to add to a Computer.

        .. deprecated:: 1.18.0

           Use :func:`add_tasks` instead.

        Parameters
        ----------
        *inputs : `.Key` or str or hashable, optional
            If provided, overrides the :attr:`inputs` property of the class.

        Returns
        -------
        tuple
            - The first, callable element of the task is :meth:`save`.
            - The second element is ``"config"``, to access the configuration of the
              Computer.
            - The third and following elements are the `inputs`.
        """
        inputs_repr = ",".join(map(repr, inputs))
        warn(
            f"Plot.make_task(…). Use: Computer.add(…, {cls.__name__}"
            + (", " if inputs_repr else "")
            + f"{inputs_repr})",
            DeprecationWarning,
        )
        return tuple([cls().save, "config"] + (list(inputs) if inputs else cls.inputs))




    @classmethod
    def add_tasks(
        cls, c: Computer, key: KeyLike, *inputs, strict: bool = False
    ) -> KeyLike:
        """Add a task to `c` to generate and save the Plot.

        Analogous to :meth:`.Operator.add_tasks`.
        """
        _inputs = list(inputs if inputs else cls.inputs)
        if strict:
            _inputs = c.check_keys(*_inputs)
        return c.add_single(key, cls().save, "config", *_inputs)




    @abstractmethod
    def generate(self, *args, **kwargs):
        """Generate and return the plot.

        A subclass of Plot **must** implement this method.

        Parameters
        ----------
        args : sequence of pandas.DataFrame or other
            One argument is given corresponding to each of the :attr:`inputs`.

            Because :mod:`plotnine` operates on pandas data structures, :meth:`save`
            automatically converts any :class:`.Quantity` inputs to
            :class:`pandas.DataFrame` before they are passed to :meth:`generate`.
        """