Add deepcopy and pickle support to figures and graph objects (#1191)

* Allow kwargs to override properties in codegen validator subclasses

This will make it easier to implement generic validator deepcopy support

* Make all validators support deepcopy

* Avoid infinite loops when checking for existence of _subplotid_props in Layout

* Add __reduce__ methods to BaseFigure and BasePlotlyType

This method provides pickle and deep copy support

* Add tests for deepcopy and pickle support

Author: jonmmease

Diff for: _plotly_utils/basevalidators.py

@@ -3,6 +3,7 @@
 import textwrap
 import uuid
 from importlib import import_module
+import copy
 
 import io
 from copy import deepcopy
@@ -373,6 +374,7 @@ def __init__(self,
         self.array_ok = array_ok
         # coerce_number is rarely used and not implemented
         self.coerce_number = coerce_number
+        self.kwargs = kwargs
 
         # Handle regular expressions
         # --------------------------
@@ -398,6 +400,17 @@ def __init__(self,
                 self.val_regexs.append(None)
                 self.regex_replacements.append(None)
 
+    def __deepcopy__(self, memodict={}):
+        """
+        A custom deepcopy method is needed here because compiled regex
+        objects don't support deepcopy
+        """
+        cls = self.__class__
+        return cls(
+            self.plotly_name,
+            self.parent_name,
+            values=self.values)
+
     @staticmethod
     def build_regex_replacement(regex_str):
         # Example: regex_str == r"^y([2-9]|[1-9][0-9]+)?$"

Diff for: codegen/validators.py

@@ -63,7 +63,7 @@ def __init__(self, plotly_name={params['plotly_name']},
             continue
 
         buffer.write(f""",
-                 {attr_name}={attr_val}""")
+                 {attr_name}=kwargs.pop('{attr_name}', {attr_val})""")
 
     buffer.write(f""",
         **kwargs""")

Diff for: plotly/animation.py

@@ -3,9 +3,10 @@
 
 class EasingValidator(EnumeratedValidator):
 
-    def __init__(self, plotly_name='easing'):
-        super(EasingValidator, self).__init__(plotly_name=plotly_name,
-                         parent_name='batch_animate',
+    def __init__(self, plotly_name='easing', parent_name='batch_animate', **_):
+        super(EasingValidator, self).__init__(
+                         plotly_name=plotly_name,
+                         parent_name=parent_name,
                          values=[
                              "linear",
                              "quad",

Diff for: plotly/basedatatypes.py

@@ -99,8 +99,9 @@ class is a subclass of both BaseFigure and widgets.DOMWidget.
         # ------------------
         # These properties are used by the tools.make_subplots logic.
         # We initialize them to None here, before checking if the input data
-        # object is a BaseFigure, in which case we bring over the _grid*
-        # properties of the input BaseFigure
+        # object is a BaseFigure, or a dict with _grid_str and _grid_ref
+        # properties, in which case we bring over the _grid* properties of
+        # the input
         self._grid_str = None
         self._grid_ref = None
 
@@ -116,6 +117,12 @@ class is a subclass of both BaseFigure and widgets.DOMWidget.
 
         elif (isinstance(data, dict)
               and ('data' in data or 'layout' in data or 'frames' in data)):
+
+            # Bring over subplot fields
+            self._grid_str = data.get('_grid_str', None)
+            self._grid_ref = data.get('_grid_ref', None)
+
+            # Extract data, layout, and frames
             data, layout, frames = (data.get('data', None),
                                     data.get('layout', None),
                                     data.get('frames', None))
@@ -230,6 +237,17 @@ class is a subclass of both BaseFigure and widgets.DOMWidget.
 
     # Magic Methods
     # -------------
+    def __reduce__(self):
+        """
+        Custom implementation of reduce is used to support deep copying
+        and pickling
+        """
+        props = self.to_dict()
+        props['_grid_str'] = self._grid_str
+        props['_grid_ref'] = self._grid_ref
+        return (self.__class__,
+                (props,))
+
     def __setitem__(self, prop, value):
 
         # Normalize prop
@@ -2594,6 +2612,15 @@ def figure(self):
 
     # Magic Methods
     # -------------
+    def __reduce__(self):
+        """
+        Custom implementation of reduce is used to support deep copying
+        and pickling
+        """
+        props = self.to_plotly_json()
+        return (self.__class__,
+                (props,))
+
     def __getitem__(self, prop):
         """
         Get item or nested item from object
@@ -3623,7 +3650,7 @@ def __getattr__(self, prop):
         Custom __getattr__ that handles dynamic subplot properties
         """
         prop = self._strip_subplot_suffix_of_1(prop)
-        if prop in self._subplotid_props:
+        if prop != '_subplotid_props' and prop in self._subplotid_props:
             validator = self._validators[prop]
             return validator.present(self._compound_props[prop])
         else:

Diff for: plotly/tests/test_io/test_deepcopy_pickle.py

@@ -0,0 +1,120 @@
+import pytest
+import copy
+import pickle
+
+from plotly.tools import make_subplots
+import plotly.graph_objs as go
+import plotly.io as pio
+
+
+# fixtures
+# --------
+@pytest.fixture
+def fig1(request):
+    return go.Figure(data=[{'type': 'scattergl',
+                            'marker': {'color': 'green'}},
+                           {'type': 'parcoords',
+                            'dimensions': [{'values': [1, 2, 3]},
+                                           {'values': [3, 2, 1]}],
+                            'line': {'color': 'blue'}}],
+                     layout={'title': 'Figure title'})
+
+
+@pytest.fixture
+def fig_subplots(request):
+    fig = make_subplots(3, 2)
+    fig.add_scatter(y=[2, 1, 3], row=1, col=1)
+    fig.add_scatter(y=[1, 3, 3], row=2, col=2)
+    return fig
+
+
+# Deep copy
+# ---------
+def test_deepcopy_figure(fig1):
+    fig_copied = copy.deepcopy(fig1)
+
+    # Contents should be equal
+    assert pio.to_json(fig_copied) == pio.to_json(fig1)
+
+    # Identities should be distinct
+    assert fig_copied is not fig1
+    assert fig_copied.layout is not fig1.layout
+    assert fig_copied.data is not fig1.data
+
+
+def test_deepcopy_figure_subplots(fig_subplots):
+    fig_copied = copy.deepcopy(fig_subplots)
+
+    # Contents should be equal
+    assert pio.to_json(fig_copied) == pio.to_json(fig_subplots)
+
+    # Subplot metadata should be equal
+    assert fig_subplots._grid_ref == fig_copied._grid_ref
+    assert fig_subplots._grid_str == fig_copied._grid_str
+
+    # Identities should be distinct
+    assert fig_copied is not fig_subplots
+    assert fig_copied.layout is not fig_subplots.layout
+    assert fig_copied.data is not fig_subplots.data
+
+    # Should be possible to add new trace to subplot location
+    fig_subplots.add_bar(y=[0, 0, 1], row=1, col=2)
+    fig_copied.add_bar(y=[0, 0, 1], row=1, col=2)
+
+    # And contents should be still equal
+    assert pio.to_json(fig_copied) == pio.to_json(fig_subplots)
+
+
+def test_deepcopy_layout(fig1):
+    copied_layout = copy.deepcopy(fig1.layout)
+
+    # Contents should be equal
+    assert copied_layout == fig1.layout
+
+    # Identities should not
+    assert copied_layout is not fig1.layout
+
+    # Original layout should still have fig1 as parent
+    assert fig1.layout.parent is fig1
+
+    # Copied layout should have no parent
+    assert copied_layout.parent is None
+
+
+# Pickling
+# --------
+def test_pickle_figure_round_trip(fig1):
+    fig_copied = pickle.loads(pickle.dumps(fig1))
+
+    # Contents should be equal
+    assert pio.to_json(fig_copied) == pio.to_json(fig1)
+
+
+def test_pickle_figure_subplots_round_trip(fig_subplots):
+    fig_copied = pickle.loads(pickle.dumps(fig_subplots))
+
+    # Contents should be equal
+    assert pio.to_json(fig_copied) == pio.to_json(fig_subplots)
+
+    # Should be possible to add new trace to subplot location
+    fig_subplots.add_bar(y=[0, 0, 1], row=1, col=2)
+    fig_copied.add_bar(y=[0, 0, 1], row=1, col=2)
+
+    # And contents should be still equal
+    assert pio.to_json(fig_copied) == pio.to_json(fig_subplots)
+
+
+def test_pickle_layout(fig1):
+    copied_layout = pickle.loads(pickle.dumps(fig1.layout))
+
+    # Contents should be equal
+    assert copied_layout == fig1.layout
+
+    # Identities should not
+    assert copied_layout is not fig1.layout
+
+    # Original layout should still have fig1 as parent
+    assert fig1.layout.parent is fig1
+
+    # Copied layout should have no parent
+    assert copied_layout.parent is None

Diff for: plotly/validators/_area.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='area', parent_name='', **kwargs):
         super(AreaValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Area',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Area'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             customdata
                 Assigns extra data each datum. This may be
                 useful when listening to hover, click and
@@ -83,6 +84,7 @@ def __init__(self, plotly_name='area', parent_name='', **kwargs):
                 visible. If "legendonly", the trace is not
                 drawn, but can appear as a legend item
                 (provided that the legend itself is visible).
-""",
+"""
+            ),
             **kwargs
         )

Diff for: plotly/validators/_bar.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='bar', parent_name='', **kwargs):
         super(BarValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Bar',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Bar'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             base
                 Sets where the bar base is drawn (in position
                 axis units). In "stack" or "relative" barmode,
@@ -210,6 +211,7 @@ def __init__(self, plotly_name='bar', parent_name='', **kwargs):
                 data.
             ysrc
                 Sets the source reference on plot.ly for  y .
-""",
+"""
+            ),
             **kwargs
         )

Diff for: plotly/validators/_box.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='box', parent_name='', **kwargs):
         super(BoxValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Box',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Box'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             boxmean
                 If True, the mean of the box(es)' underlying
                 distribution is drawn as a dashed line inside
@@ -179,6 +180,7 @@ def __init__(self, plotly_name='box', parent_name='', **kwargs):
                 data.
             ysrc
                 Sets the source reference on plot.ly for  y .
-""",
+"""
+            ),
             **kwargs
         )

Diff for: plotly/validators/_candlestick.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='candlestick', parent_name='', **kwargs):
         super(CandlestickValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Candlestick',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Candlestick'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             close
                 Sets the close values.
             closesrc
@@ -129,6 +130,7 @@ def __init__(self, plotly_name='candlestick', parent_name='', **kwargs):
                 (the default value), the y coordinates refer to
                 `layout.yaxis`. If "y2", the y coordinates
                 refer to `layout.yaxis2`, and so on.
-""",
+"""
+            ),
             **kwargs
         )

Diff for: plotly/validators/_carpet.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='carpet', parent_name='', **kwargs):
         super(CarpetValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Carpet',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Carpet'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             a
                 An array containing values of the first
                 parameter value
@@ -139,6 +140,7 @@ def __init__(self, plotly_name='carpet', parent_name='', **kwargs):
                 refer to `layout.yaxis2`, and so on.
             ysrc
                 Sets the source reference on plot.ly for  y .
-""",
+"""
+            ),
             **kwargs
         )

Diff for: plotly/validators/_choropleth.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='choropleth', parent_name='', **kwargs):
         super(ChoroplethValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Choropleth',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Choropleth'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             autocolorscale
                 Determines whether the colorscale is a default
                 palette (`autocolorscale: true`) or the palette
@@ -150,6 +151,7 @@ def __init__(self, plotly_name='choropleth', parent_name='', **kwargs):
                 set, `zmax` must be set as well.
             zsrc
                 Sets the source reference on plot.ly for  z .
-""",
+"""
+            ),
             **kwargs
         )

Diff for: plotly/validators/_cone.py

@@ -7,8 +7,9 @@ def __init__(self, plotly_name='cone', parent_name='', **kwargs):
         super(ConeValidator, self).__init__(
             plotly_name=plotly_name,
             parent_name=parent_name,
-            data_class_str='Cone',
-            data_docs="""
+            data_class_str=kwargs.pop('data_class_str', 'Cone'),
+            data_docs=kwargs.pop(
+                'data_docs', """
             anchor
                 Sets the cones' anchor with respect to their
                 x/y/z positions. Note that "cm" denote the
@@ -189,6 +190,7 @@ def __init__(self, plotly_name='cone', parent_name='', **kwargs):
                 of the displayed cones.
             zsrc
                 Sets the source reference on plot.ly for  z .
-""",
+"""
+            ),
             **kwargs
         )