Handle groups of plots

But no difference between grouped plots and ungrouped plots at the
moment.

replot/__init__.py

@@ -2,6 +2,7 @@
 The :mod:`replot` module is a (sane) Python plotting module, abstracting on top
 of Matplotlib.
 """
+import collections
 import shutil
 
 import matplotlib.pyplot as plt
@@ -15,25 +16,6 @@ from replot import exceptions as exc
 __VERSION__ = "0.0.1"
 
 
-def mpl_custom_rc_context():
-    """
-    Overload ``matplotlib.rcParams`` to enable advanced features if \
-            available. In particular, use LaTeX if available.
-
-    :returns: A ``matplotlib.rc_context`` object to use in a ``with`` \
-            statement.
-    """
-    custom_rc = {}
-    # Add LaTeX in rc if available
-    if(shutil.which("latex") is not None and
-       shutil.which("gs") is not None and
-       shutil.which("dvipng") is not None):
-        # LateX dependencies are all available
-        custom_rc["text.usetex"] = True
-        custom_rc["text.latex.unicode"] = True
-    return plt.rc_context(rc=custom_rc)
-
-
 class Figure():
     """
     The main class from :mod:`replot`, representing a figure. Can be used \
@@ -69,7 +51,7 @@ class Figure():
         self.palette = palette
         self.max_colors = max_colors
         self.legend = legend
-        self.plots = []
+        self.plots = collections.defaultdict(list)  # keys are groups
 
     def __enter__(self):
         return self
@@ -91,12 +73,13 @@ class Figure():
                 # Create figure
                 figure, axes = plt.subplots()
                 # Add plots
-                for plot_ in self.plots:
+                for group_ in self.plots:
-                    tmp_plots = axes.plot(*(plot_[0]), **(plot_[1]))
+                    for plot_ in self.plots[group_]:
-                    # Do not clip line at the axes boundaries to prevent
+                        tmp_plots = axes.plot(*(plot_[0]), **(plot_[1]))
-                    # extremas from being cropped.
+                        # Do not clip line at the axes boundaries to prevent
-                    for tmp_plot in tmp_plots:
+                        # extremas from being cropped.
-                        tmp_plot.set_clip_on(False)
+                        for tmp_plot in tmp_plots:
+                            tmp_plot.set_clip_on(False)
                 # Set properties
                 axes.set_xlabel(self.xlabel)
                 axes.set_ylabel(self.ylabel)
@@ -123,66 +106,49 @@ class Figure():
             - Two arguments being a function and an interval represented by \
                     a tuple of its bounds.
 
-        .. note:: ``kwargs`` arguments are directly passed to \
+        .. note:: ``kwargs`` arguments are passed to \
                     ``matplotlib.pyplot.plot``.
 
+        .. note:: You can use some :mod:`replot` specific keyword arguments:
+
+            - ``group`` which defines permits to group plots together, in \
+                    subplots (one unicode character maximum).
+
         >>> with replot.figure() as fig: fig.plot(np.sin, (-1, 1))
         >>> with replot.figure() as fig: fig.plot(np.sin, [-1, -0.9, …, 1])
         >>> with replot.figure() as fig: fig.plot([1, 2, 3], [4, 5, 6])
         >>> with replot.figure() as fig: fig.plot([1, 2, 3],
                                                   [4, 5, 6], linewidth=2.0)
+        >>> with replot.figure() as fig: fig.plot([1, 2, 3],
+                                                  [4, 5, 6], group="a")
         """
         if len(args) == 0:
             raise exc.InvalidParameterError(
                 "You should pass at least one argument to this function.")
 
+        # Extra custom kwargs (the ones from replot but not matplotlib) from
+        # kwargs
+        kwargs, custom_kwargs = _handle_custom_plot_arguments(kwargs)
+
         if hasattr(args[0], "__call__"):
             # We want to plot a function
-            self._plot_function(args[0], *(args[1:]), **kwargs)
+            plot_ = _plot_function(args[0], *(args[1:]), **kwargs)
         else:
             # Else, it is a point series, and we just have to store it for
             # later plotting.
-            self.plots.append((args, kwargs))
+            plot_ = (args, kwargs)
+
+        # Add the plot to the correct group
+        if "group" in custom_kwargs:
+            self.plots[custom_kwargs["group"]].append(plot_)
+        else:
+            self.plots["default"].append(plot_)
 
         # Automatically set the legend if label is found
         # (only do it if legend is not explicitly suppressed)
         if "label" in kwargs and self.legend is None:
             self.legend = True
 
-    def _plot_function(self, data, *args, **kwargs):
-        """
-        Helper function to handle plotting of unevaluated functions (trying \
-                to evaluate it nicely and rendering the plot).
-
-        :param data: The function to plot.
-
-        .. seealso:: The documentation of the ``replot.Figure.plot`` method.
-
-        .. note:: ``args`` is used to handle the interval or point series on \
-                which the function should be evaluated. ``kwargs`` are passed \
-                directly to ``matplotlib.pyplot.plot`.
-        """
-        if len(args) == 0:
-            # If no interval specified, raise an issue
-            raise exc.InvalidParameterError(
-                "You should pass a plotting interval to the plot command.")
-        elif isinstance(args[0], tuple):
-            # Interval specified, use it and adaptive plotting
-            x_values, y_values = adaptive_sampling.sample_function(
-                data,
-                args[0],
-                tol=1e-3)
-        elif isinstance(args[0], (list, np.ndarray)):
-            # List of points specified, use them and compute values of the
-            # function
-            x_values = args[0]
-            y_values = [data(i) for i in x_values]
-        else:
-            raise exc.InvalidParameterError(
-                "Second parameter in plot command should be a tuple " +
-                "specifying plotting interval.")
-        self.plots.append(((x_values, y_values) + args[1:], kwargs))
-
     def _legend(self, axes):
         """
         Helper function to handle ``legend`` attribute. It places the legend \
@@ -250,3 +216,79 @@ def plot(data, **kwargs):
         else:
             figure.plot(plot_)
     figure.show()
+
+
+def mpl_custom_rc_context():
+    """
+    Overload ``matplotlib.rcParams`` to enable advanced features if \
+            available. In particular, use LaTeX if available.
+
+    :returns: A ``matplotlib.rc_context`` object to use in a ``with`` \
+            statement.
+    """
+    custom_rc = {}
+    # Add LaTeX in rc if available
+    if(shutil.which("latex") is not None and
+       shutil.which("gs") is not None and
+       shutil.which("dvipng") is not None):
+        # LateX dependencies are all available
+        custom_rc["text.usetex"] = True
+        custom_rc["text.latex.unicode"] = True
+    return plt.rc_context(rc=custom_rc)
+
+
+def _handle_custom_plot_arguments(kwargs):
+    """
+    This method handles custom keyword arguments from plot in \
+            :mod:`replot` which are not in :mod:`matplotlib` function.
+
+    :param kwargs: A dictionary of keyword arguments to handle.
+    :return: A tuple of :mod:`matplotlib` compatible keyword arguments \
+            and of extra :mod:`replot` keyword arguments, both returned \
+            as ``dict``.
+    """
+    custom_kwargs = {}
+    # Handle "group" argument
+    if "group" in kwargs:
+        if len(kwargs["group"]) > 1:
+            raise exc.InvalidParameterError(
+                "Group name cannot be longer than one unicode character.")
+        custom_kwargs["group"] = kwargs["group"]
+        del kwargs["group"]
+        return (kwargs, custom_kwargs)
+
+
+def _plot_function(data, *args, **kwargs):
+    """
+    Helper function to handle plotting of unevaluated functions (trying \
+            to evaluate it nicely and rendering the plot).
+
+    :param data: The function to plot.
+    :returns: A tuple of ``(args, kwargs)`` representing the plot.
+
+    .. seealso:: The documentation of the ``replot.Figure.plot`` method.
+
+    .. note:: ``args`` is used to handle the interval or point series on \
+            which the function should be evaluated. ``kwargs`` are passed \
+            directly to ``matplotlib.pyplot.plot`.
+    """
+    if len(args) == 0:
+        # If no interval specified, raise an issue
+        raise exc.InvalidParameterError(
+            "You should pass a plotting interval to the plot command.")
+    elif isinstance(args[0], tuple):
+        # Interval specified, use it and adaptive plotting
+        x_values, y_values = adaptive_sampling.sample_function(
+            data,
+            args[0],
+            tol=1e-3)
+    elif isinstance(args[0], (list, np.ndarray)):
+        # List of points specified, use them and compute values of the
+        # function
+        x_values = args[0]
+        y_values = [data(i) for i in x_values]
+    else:
+        raise exc.InvalidParameterError(
+            "Second parameter in plot command should be a tuple " +
+            "specifying plotting interval.")
+    return ((x_values, y_values) + args[1:], kwargs)