Added dict-style pop method to Figure and graph objects

Author: jonmmease

Diff for: packages/python/plotly/plotly/basedatatypes.py

@@ -498,6 +498,37 @@ def update(self, dict1=None, **kwargs):
                             self[k] = v
         return self
 
+    def pop(self, key, *args):
+        """
+        Remove the value associated with the specified key and return it
+
+        Parameters
+        ----------
+        key: str
+            Property name
+        dflt
+            The default value to return if key was not found in figure
+
+        Returns
+        -------
+        value
+            The removed value that was previously associated with key
+
+        Raises
+        ------
+        KeyError
+            If key is not in object and no dflt argument specified
+        """
+        # Handle default
+        if key not in self and args:
+            return args[0]
+        elif key in self:
+            val = self[key]
+            self[key] = None
+            return val
+        else:
+            raise KeyError(key)
+
     # Data
     # ----
     @property
@@ -520,6 +551,10 @@ def data(self, new_data):
                       'a list or tuple that contains a permutation of a '
                       'subset of itself.\n')
 
+        # ### Treat None as empty ###
+        if new_data is None:
+            new_data = ()
+
         # ### Check valid input type ###
         if not isinstance(new_data, (list, tuple)):
             err_msg = (err_header + '    Received value with type {typ}'
@@ -3503,6 +3538,37 @@ def update(self, dict1=None, **kwargs):
 
         return self
 
+    def pop(self, key, *args):
+        """
+        Remove the value associated with the specified key and return it
+
+        Parameters
+        ----------
+        key: str
+            Property name
+        dflt
+            The default value to return if key was not found in object
+
+        Returns
+        -------
+        value
+            The removed value that was previously associated with key
+
+        Raises
+        ------
+        KeyError
+            If key is not in object and no dflt argument specified
+        """
+        # Handle default
+        if key not in self and args:
+            return args[0]
+        elif key in self:
+            val = self[key]
+            self[key] = None
+            return val
+        else:
+            raise KeyError(key)
+
     @property
     def _in_batch_mode(self):
         """

Diff for: packages/python/plotly/plotly/tests/test_core/test_graph_objs/test_figure.py

@@ -146,3 +146,18 @@ def test_scalar_trace_as_data(self):
 
         fig = go.Figure(data=dict(type='waterfall', y=[2, 1, 3]))
         self.assertEqual(fig.data, (go.Waterfall(y=[2, 1, 3]),))
+
+    def test_pop_data(self):
+        fig = go.Figure(data=go.Waterfall(y=[2, 1, 3]))
+        self.assertEqual(fig.pop('data'), (go.Waterfall(y=[2, 1, 3]),))
+        self.assertEqual(fig.data, ())
+
+    def test_pop_layout(self):
+        fig = go.Figure(layout=go.Layout(width=1000))
+        self.assertEqual(fig.pop('layout'), go.Layout(width=1000))
+        self.assertEqual(fig.layout, go.Layout())
+
+    def test_pop_invalid_key(self):
+        fig = go.Figure(layout=go.Layout(width=1000))
+        with self.assertRaises(KeyError):
+            fig.pop('bogus')

Diff for: packages/python/plotly/plotly/tests/test_core/test_graph_objs/test_graph_objs.py

@@ -115,3 +115,40 @@ def test_legacy_title_props_remapped(self):
         self.assertIn('titlefont', obj)
         self.assertIn('titlefont.family', obj)
         self.assertIn('titlefont', iter(obj))
+
+
+class TestPop(TestCase):
+    def setUp(self):
+        self.layout = go.Layout(
+            width=1000,
+            title={'text': 'the title', 'font': {'size': 20}},
+            annotations=[{}, {}],
+            xaxis2={'range': [1, 2]}
+        )
+
+    def test_pop_valid_simple_prop(self):
+        self.assertEqual(self.layout.width, 1000)
+        self.assertEqual(self.layout.pop('width'), 1000)
+        self.assertIsNone(self.layout.width)
+
+    def test_pop_valid_compound_prop(self):
+        val = self.layout.title
+        self.assertEqual(self.layout.pop('title'), val)
+        self.assertEqual(self.layout.title, go.layout.Title())
+
+    def test_pop_valid_array_prop(self):
+        val = self.layout.annotations
+        self.assertEqual(self.layout.pop('annotations'), val)
+        self.assertEqual(self.layout.annotations, ())
+
+    def test_pop_valid_subplot_prop(self):
+        val = self.layout.xaxis2
+        self.assertEqual(self.layout.pop('xaxis2'), val)
+        self.assertEqual(self.layout.xaxis2, go.layout.XAxis())
+
+    def test_pop_invalid_prop_key_error(self):
+        with self.assertRaises(KeyError):
+            self.layout.pop('bogus')
+
+    def test_pop_invalid_prop_with_default(self):
+        self.assertEqual(self.layout.pop('bogus', 42), 42)