Merge pull request #2866 from plotly/fig-set-subplots

`go.Figure.set_subplots`

Author: nicholas-esterer

Diff for: CHANGELOG.md

@@ -4,6 +4,10 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [4.13.0] - UNRELEASED
 
+### Added
+
+- `go.Figure` now has a `set_subplots` method to set subplots on an already
+  existing figure.
 
 
 ## [4.12.1] - UNRELEASED

Diff for: doc/python/subplots.md

@@ -581,6 +581,26 @@ fig = go.Figure(data=data, layout=layout)
 fig.show()
 ```
 
+#### Setting Subplots on a Figure Directly
+
+_new in 4.13_
+
+Subplots can be added to an already existing figure, provided it doesn't already
+have subplots. `go.Figure.set_subplots` accepts all the same arguments as
+`plotly.subplots.make_subplots`.
+
+```python
+import plotly.graph_objects as go
+fig = go.Figure().set_subplots(2, 3, horizontal_spacing=0.1)
+```
+
+is equivalent to:
+
+```python
+from plotly.subplots import make_subplots
+fig = make_subplots(2, 3, horizontal_spacing=0.1)
+```
+
 #### Reference
 All of the x-axis properties are found here: https://plotly.com/python/reference/XAxis/
 All of the y-axis properties are found here: https://plotly.com/python/reference/YAxis/

Diff for: 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
@@ -3935,6 +3936,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 for: 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):

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

@@ -1,7 +1,9 @@
 from __future__ import absolute_import
 
 import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 from plotly.tests.utils import TestCaseNoTemplate
+import pytest
 
 
 class FigureTest(TestCaseNoTemplate):
@@ -199,3 +201,21 @@ def test_update_overwrite_data(self):
             fig.to_plotly_json()["data"],
             [{"type": "scatter", "y": [1, 3, 2], "line": {"color": "yellow"}}],
         )
+
+
+def test_set_subplots():
+    # Test that it works the same as make_subplots for a simple call
+    fig0 = go.Figure()
+    fig0_sp = make_subplots(2, 2)
+    fig0.set_subplots(2, 2)
+    assert fig0.layout == fig0_sp.layout
+    # Test that it accepts the same arguments as make_subplots
+    fig1 = go.Figure()
+    fig1.set_subplots(rows=2, cols=2, horizontal_spacing=0.25, vertical_spacing=0.1)
+    fig1_sp = make_subplots(
+        rows=2, cols=2, horizontal_spacing=0.25, vertical_spacing=0.1
+    )
+    assert fig1.layout == fig1_sp.layout
+    # Test that calling on a figure that already has subplots throws an error.
+    with pytest.raises(ValueError, match=r"^This figure already has subplots\.$"):
+        fig1.set_subplots(2, 3)

Diff for: packages/python/plotly/plotly/tests/test_core/test_subplots/test_make_subplots.py

@@ -1924,3 +1924,13 @@ def test_secondary_y_subplots(self):
         expected.update_traces(uid=None)
 
         self.assertEqual(fig.to_plotly_json(), expected.to_plotly_json())
+
+    def test_if_passed_figure(self):
+        # assert it returns the same figure it was passed
+        fig = Figure()
+        figsp = subplots.make_subplots(2, 2, figure=fig)
+        assert id(fig) == id(figsp)
+        # assert the layout is the same when it returns its own figure
+        fig2sp = subplots.make_subplots(2, 2)
+        assert id(fig2sp) != id(figsp)
+        assert fig2sp.layout == figsp.layout