go.Figure now has set_subplots

nicholas-esterer · nicholas-esterer · commit 9d350bc880a8 · 2020-10-30T12:06:47.000-04:00
Calling `fig.set_subplots(2,3)` works just like fig=make_subplots(2,3).
Also accepts the same keywords arguments as make_subplots. Fails if fig
already has subplots.

Tests still need to be added.
diff --git a/packages/python/plotly/plotly/basedatatypes.py b/packages/python/plotly/plotly/basedatatypes.py
@@ -14,6 +14,7 @@
 from .optional_imports import get_module
 
 from . import shapeannotation
+from . import subplots
 
 # Create Undefined sentinel value
 #   - Setting a property to None removes any existing value
@@ -3926,6 +3927,22 @@ def _subplot_not_empty(self, xref, yref, selector="all"):
                 )
         return ret
 
+    def set_subplots(self, rows=None, cols=None, **make_subplots_args):
+        """
+        Add subplots to this figure. If the figure already contains subplots,
+        then this throws an error. Accepts any keyword arguments that
+        plotly.subplots.make_subplots accepts.
+        """
+        # rows, cols provided so that this can be called like
+        # fig.set_subplots(2,3), say
+        if rows is not None:
+            make_subplots_args["rows"] = rows
+        if cols is not None:
+            make_subplots_args["cols"] = cols
+        if self._has_subplots():
+            raise ValueError("This figure already has subplots.")
+        return subplots.make_subplots(figure=self, **make_subplots_args)
+
 
 class BasePlotlyType(object):
     """
diff --git a/packages/python/plotly/plotly/subplots.py b/packages/python/plotly/plotly/subplots.py
@@ -60,6 +60,7 @@ def make_subplots(
     row_titles=None,
     x_title=None,
     y_title=None,
+    figure=None,
     **kwargs
 ):
     """
@@ -226,7 +227,15 @@ def make_subplots(
 
     y_title: str or None (default None)
         Title to place to the left of the left column of subplots,
-       centered vertically
+        centered vertically
+
+    figure: go.Figure or None (default None)
+        If None, a new go.Figure instance will be created and its axes will be
+        populated with those corresponding to the requested subplot geometry and
+        this new figure will be returned.
+        If a go.Figure instance, the axes will be added to the
+        layout of this figure and this figure will be returned. If the figure
+        already contains axes, they will be overwritten.
 
     Examples
     --------
@@ -809,13 +818,15 @@ def _checks(item, defaults):
         print(grid_str)
 
     # Build resulting figure
-    fig = go.Figure(layout=layout)
+    if figure is None:
+        figure = go.Figure()
+    figure.update_layout(layout)
 
     # Attach subplot grid info to the figure
-    fig.__dict__["_grid_ref"] = grid_ref
-    fig.__dict__["_grid_str"] = grid_str
+    figure.__dict__["_grid_ref"] = grid_ref
+    figure.__dict__["_grid_str"] = grid_str
 
-    return fig
+    return figure
 
 
 def _configure_shared_axes(layout, grid_ref, specs, x_or_y, shared, row_dir):
