V4 Update add_trace, add_traces, and add_{trace} methods to return figure

packages/python/chart-studio/chart_studio/plotly/plotly.py

@@ -1793,7 +1793,7 @@ def _extract_grid_from_fig_like(fig, grid=None, path=''):
         trace_type = trace_dict.get('type', 'scatter')
         if trace_type not in reference_traces:
             reference_traces[trace_type] = reference_fig.add_trace(
-                {'type': trace_type})
+                {'type': trace_type}).data[-1]
 
         reference_trace = reference_traces[trace_type]
         _extract_grid_graph_obj(

packages/python/plotly/codegen/datatypes.py

@@ -265,7 +265,13 @@ def __init__(self""")
                f'dict of properties compatible with this constructor '
                f'or an instance of {class_name}')]
 
-    add_docstring(buffer, node, header=header, prepend_extras=extras)
+    add_docstring(
+        buffer,
+        node,
+        header=header,
+        prepend_extras=extras,
+        return_type=node.name_datatype_class,
+    )
 
     buffer.write(f"""
         super({datatype_class}, self).__init__('{node.name_property}')
@@ -424,7 +430,14 @@ def add_constructor_params(buffer,
         ):""")
 
 
-def add_docstring(buffer, node, header, prepend_extras=(), append_extras=()):
+def add_docstring(
+        buffer,
+        node,
+        header,
+        prepend_extras=(),
+        append_extras=(),
+        return_type=None,
+):
     """
     Write docstring for a compound datatype node
 
@@ -443,6 +456,8 @@ def add_docstring(buffer, node, header, prepend_extras=(), append_extras=()):
     append_extras :
         List or tuple of propery name / description pairs that should be
         included at the end of the docstring
+    return_type :
+        The docstring return type
     Returns
     -------
 
@@ -511,7 +526,7 @@ def add_docstring(buffer, node, header, prepend_extras=(), append_extras=()):
 
         Returns
         -------
-        {node.name_datatype_class}
+        {return_type}
         \"\"\"""")
 
 

packages/python/plotly/codegen/figure.py

@@ -110,26 +110,6 @@ def __init__(self, data=None, layout=None,
             d for d in trace_node.child_datatypes
             if d.name_property == 'yaxis'
         ])
-        if include_secondary_y:
-            secondary_y_1 = ', secondary_y=None'
-            secondary_y_2 = ', secondary_y=secondary_y'
-            secondary_y_docstring = f"""
-        secondary_y: boolean or None (default None)
-            * If True, only select yaxis objects associated with the secondary
-              y-axis of the subplot.
-            * If False, only select yaxis objects associated with the primary
-              y-axis of the subplot.
-            * If None (the default), do not filter yaxis objects based on
-              a secondary y-axis condition. 
-
-            To select yaxis objects by secondary y-axis, the Figure must
-            have been created using plotly.subplots.make_subplots. See
-            the docstring for the specs argument to make_subplots for more
-            info on creating subplots with secondary y-axes."""
-        else:
-            secondary_y_1 = ''
-            secondary_y_2 = ''
-            secondary_y_docstring = ''
 
         # #### Function signature ####
         buffer.write(f"""
@@ -171,7 +151,13 @@ def add_{trace_node.plotly_name}(self""")
 """)
             )
 
-        add_docstring(buffer, trace_node, header, append_extras=doc_extras)
+        add_docstring(
+            buffer,
+            trace_node,
+            header,
+            append_extras=doc_extras,
+            return_type=fig_classname,
+        )
 
         # #### Function body ####
         buffer.write(f"""

packages/python/plotly/plotly/basedatatypes.py

@@ -1447,12 +1447,12 @@ def add_trace(self, trace, row=None, col=None, secondary_y=None):
                 (scatter, bar, etc.)
         Returns
         -------
-        BaseTraceType
-            The newly added trace
+        BaseFigure
+            The Figure that add_trace was called on
 
         Examples
         --------
-        >>> from plotly import tools
+        >>> from plotly import subplots
         >>> import plotly.graph_objs as go
 
         Add two Scatter traces to a figure
@@ -1462,11 +1462,7 @@ def add_trace(self, trace, row=None, col=None, secondary_y=None):
 
 
         Add two Scatter traces to vertically stacked subplots
-        >>> fig = tools.make_subplots(rows=2)
-        This is the format of your plot grid:
-        [ (1,1) x1,y1 ]
-        [ (2,1) x2,y2 ]
-
+        >>> fig = subplots.make_subplots(rows=2)
         >>> fig.add_trace(go.Scatter(x=[1,2,3], y=[2,1,2]), row=1, col=1)
         >>> fig.add_trace(go.Scatter(x=[1,2,3], y=[2,1,2]), row=2, col=1)
         """
@@ -1492,7 +1488,7 @@ def add_trace(self, trace, row=None, col=None, secondary_y=None):
             rows=[row] if row is not None else None,
             cols=[col] if col is not None else None,
             secondary_ys=[secondary_y] if secondary_y is not None else None
-        )[0]
+        )
 
     def add_traces(self, data, rows=None, cols=None, secondary_ys=None):
         """
@@ -1528,12 +1524,12 @@ def add_traces(self, data, rows=None, cols=None, secondary_ys=None):
 
         Returns
         -------
-        tuple[BaseTraceType]
-            Tuple of the newly added traces
+        BaseFigure
+            The Figure that add_traces was called on
 
         Examples
         --------
-        >>> from plotly import tools
+        >>> from plotly import subplots
         >>> import plotly.graph_objs as go
 
         Add two Scatter traces to a figure
@@ -1542,11 +1538,7 @@ def add_traces(self, data, rows=None, cols=None, secondary_ys=None):
         ...                 go.Scatter(x=[1,2,3], y=[2,1,2])])
 
         Add two Scatter traces to vertically stacked subplots
-        >>> fig = tools.make_subplots(rows=2)
-        This is the format of your plot grid:
-        [ (1,1) x1,y1 ]
-        [ (2,1) x2,y2 ]
-
+        >>> fig = subplots.make_subplots(rows=2)
         >>> fig.add_traces([go.Scatter(x=[1,2,3], y=[2,1,2]),
         ...                 go.Scatter(x=[1,2,3], y=[2,1,2])],
         ...                 rows=[1, 2], cols=[1, 1])
@@ -1608,7 +1600,7 @@ def add_traces(self, data, rows=None, cols=None, secondary_ys=None):
         # Update messages
         self._send_addTraces_msg(new_traces_data)
 
-        return data
+        return self
 
     # Subplots
     # --------