Add Figure.to_dataframe method.

theengineear · theengineear · commit e116ff2a080a · 2015-04-09T01:46:14.000-04:00
diff --git a/plotly/graph_objs/graph_objs.py b/plotly/graph_objs/graph_objs.py
@@ -1014,6 +1014,16 @@ def get_data(self, flatten=False):
             return data['data']
     Figure.get_data = get_data
 
+    def to_dataframe(self):
+        data = self.get_data(flatten=True)
+        try:
+            import pandas as pd
+            return pd.DataFrame(dict([(k, pd.Series(v))
+                                      for k, v in data.items()]))
+        except ImportError:
+            return data
+    Figure.to_dataframe = to_dataframe
+
     def append_trace(self, trace, row, col):
         """ Helper function to add a data traces to your figure
         that is bound to axes at the row, col index.
diff --git a/plotly/tests/test_optional/test_graph_objs/__init__.py b/plotly/tests/test_optional/test_graph_objs/__init__.py
diff --git a/plotly/tests/test_optional/test_graph_objs/test_to_dataframe.py b/plotly/tests/test_optional/test_graph_objs/test_to_dataframe.py
@@ -0,0 +1,107 @@
+from __future__ import absolute_import
+
+from unittest import TestCase
+
+from plotly.graph_objs import *
+
+
+class TestToDataframe(TestCase):
+
+    fig = None
+
+    def setUp(self):
+        super(TestToDataframe, self).setUp()
+        self.fig = Figure(
+            data=Data([
+                Scatter(
+                    x=[52698, 43117],
+                    y=[53, 31],
+                    mode='markers',
+                    name='North America',
+                    text=['United States', 'Canada'],
+                    marker=Marker(
+                        color='rgb(164, 194, 244)',
+                        size=12,
+                        line=Line(
+                            color='white',
+                            width=0.5
+                        )
+                    )
+                ),
+                Scatter(
+                    x=[39317, 37236, 35650, 30066, 29570, 27159, 23557, 21046,
+                       18007],
+                    y=[33, 20, 13, 19, 27, 19, 49, 44, 38],
+                    mode='markers',
+                    name='Europe',
+                    text=['Germany', 'Britain', 'France', 'Spain', 'Italy',
+                          'Czech Rep.', 'Greece', 'Poland'],
+                    marker=Marker(
+                        color='rgb(255, 217, 102)',
+                        size=12,
+                        line=Line(
+                            color='white',
+                            width=0.5
+                        )
+                    )
+                ),
+                Scatter(
+                    x=[42952, 37037, 33106, 17478, 9813, 5253, 4692, 3899],
+                    y=[23, 42, 54, 89, 14, 99, 93, 70],
+                    mode='markers',
+                    name='Asia/Pacific',
+                    text=['Australia', 'Japan', 'South Korea', 'Malaysia',
+                          'China', 'Indonesia', 'Philippines', 'India'],
+                    marker=Marker(
+                        color='rgb(234, 153, 153)',
+                        size=12,
+                        line=Line(
+                            color='white',
+                            width=0.5
+                        )
+                    )
+                ),
+                Scatter(
+                    x=[19097, 18601, 15595, 13546, 12026, 7434, 5419],
+                    y=[43, 47, 56, 80, 86, 93, 80],
+                    mode='markers',
+                    name='Latin America',
+                    text=['Chile', 'Argentina', 'Mexico', 'Venezuela',
+                          'Venezuela', 'El Salvador', 'Bolivia'],
+                    marker=Marker(
+                        color='rgb(142, 124, 195)',
+                        size=12,
+                        line=Line(
+                            color='white',
+                            width=0.5
+                        )
+                    )
+                )
+            ]),
+            layout=Layout(
+                title='Quarter 1 Growth',
+                autosize=False,
+                width=500,
+                height=500,
+                xaxis=XAxis(
+                    title='GDP per Capita',
+                    showgrid=False,
+                    zeroline=False
+                ),
+                yaxis=YAxis(
+                    title='Percent',
+                    showline=False
+                ),
+                margin=Margin(
+                    l=65,
+                    r=50,
+                    b=65,
+                    t=90
+                )
+            )
+        )
+
+    def test_figure_to_dataframe(self):
+        df = self.fig.to_dataframe()
+        self.assertEqual(len(df), 9)
+        self.assertEqual(len(df.columns), 12)
