Add select_*, for_each_*, update_* methods for each layout subplot type

Also updated code generation logic to automatically determine the set
of layout subplot types.

Author: jonmmease

codegen/__init__.py

@@ -126,6 +126,9 @@ def perform_codegen():
     all_layout_nodes = PlotlyNode.get_all_datatype_nodes(
         plotly_schema, LayoutNode)
 
+    subplot_nodes = [node for node in layout_node.child_compound_datatypes
+                     if node.node_data.get('_isSubplotObj', False)]
+
     # ### FrameNode ###
     compound_frame_nodes = PlotlyNode.get_all_compound_datatype_nodes(
         plotly_schema, FrameNode)
@@ -187,7 +190,8 @@ def perform_codegen():
                          base_traces_node,
                          data_validator,
                          layout_validator,
-                         frame_validator)
+                         frame_validator,
+                         subplot_nodes)
 
     # Write datatype __init__.py files
     # --------------------------------

codegen/datatypes.py

@@ -103,6 +103,48 @@ def build_datatype_py(node):
 
 class {datatype_class}(_{node.name_base_datatype}):\n""")
 
+    # ### Layout subplot properties ###
+    if datatype_class == 'Layout':
+        subplot_nodes = [node for node in node.child_compound_datatypes
+                         if node.node_data.get('_isSubplotObj', False)]
+        subplot_names = [n.name_property for n in subplot_nodes]
+        buffer.write(f"""
+    _subplotid_prop_names = {repr(subplot_names)}
+    
+    import re
+    _subplotid_prop_re = re.compile(
+        '^(' + '|'.join(_subplotid_prop_names) + ')(\d+)$')
+""")
+
+        subplot_validator_names = [n.name_validator_class
+                                   for n in subplot_nodes]
+
+        validator_csv = ', '.join(subplot_validator_names)
+        subplot_dict_str = (
+                '{' +
+                ', '.join(f"'{subname}': {valname}" for subname, valname in
+                          zip(subplot_names, subplot_validator_names)) +
+                '}'
+        )
+
+        buffer.write(f"""
+    @property
+    def _subplotid_validators(self):
+        \"\"\"
+        dict of validator classes for each subplot type
+
+        Returns
+        -------
+        dict
+        \"\"\"
+        from plotly.validators.layout import ({validator_csv})
+
+        return {subplot_dict_str}
+        
+    def _subplot_re_match(self, prop):
+        return self._subplotid_prop_re.match(prop)
+""")
+
     # ### Property definitions ###
     child_datatype_nodes = node.child_datatypes
 

codegen/figure.py

@@ -8,9 +8,12 @@
                                add_constructor_params, add_docstring)
 from codegen.utils import PlotlyNode, format_and_write_source_py
 
+import inflect
+
 
 def build_figure_py(trace_node, base_package, base_classname, fig_classname,
-                    data_validator, layout_validator, frame_validator):
+                    data_validator, layout_validator, frame_validator,
+                    subplot_nodes):
     """
 
     Parameters
@@ -30,7 +33,8 @@ def build_figure_py(trace_node, base_package, base_classname, fig_classname,
         LayoutValidator instance
     frame_validator : CompoundArrayValidator
         FrameValidator instance
-
+    subplot_nodes: list of str
+        List of names of all of the layout subplot properties
     Returns
     -------
     str
@@ -53,6 +57,7 @@ def build_figure_py(trace_node, base_package, base_classname, fig_classname,
     # ### Import trace graph_obj classes ###
     trace_types_csv = ', '.join([n.name_datatype_class for n in trace_nodes])
     buffer.write(f'from plotly.graph_objs import ({trace_types_csv})\n')
+    buffer.write("from plotly.subplots import _validate_v4_subplots\n")
 
     # Write class definition
     # ----------------------
@@ -141,6 +146,111 @@ def add_{trace_node.plotly_name}(self""")
         buffer.write(f"""
         return self.add_trace(new_trace, row=row, col=col)""")
 
+    # update layout subplots
+    # ----------------------
+    inflect_eng = inflect.engine()
+    for subplot_node in subplot_nodes:
+        singular_name = subplot_node.name_property
+        plural_name = inflect_eng.plural_noun(singular_name)
+        buffer.write(f"""
+
+    def select_{plural_name}(self, selector=None, row=None, col=None):
+        \"\"\"
+        Select {singular_name} subplot objects from a particular subplot cell
+        and/or {singular_name} subplot objects that satisfy custom selection
+        criteria.
+
+        Parameters
+        ----------
+        selector: dict or None (default None)
+            Dict to use as selection criteria.
+            {singular_name} objects will be selected if they contain
+            properties corresponding to all of the dictionary's keys, with
+            values that exactly match the supplied values. If None
+            (the default), all {singular_name} objects are selected.
+        row, col: int or None (default None)
+            Subplot row and column index of {singular_name} objects to select.
+            To select {singular_name} objects by row and column, the Figure
+            must have been created using plotly.subplots.make_subplots.
+            If None (the default), all {singular_name} objects are selected.
+
+        Returns
+        -------
+        generator
+            Generator that iterates through all of the {singular_name}
+            objects that satisfy all of the specified selection criteria
+        \"\"\"
+        if row is not None or col is not None:
+            _validate_v4_subplots('select_{plural_name}')
+
+        return self._select_layout_subplots_by_prefix(
+            '{singular_name}', selector, row, col)
+
+    def for_each_{singular_name}(self, fn, selector=None, row=None, col=None):
+        \"\"\"
+        Apply a function to all {singular_name} objects that satisfy the
+        specified selection criteria
+        
+        Parameters
+        ----------
+        fn:
+            Function that inputs a single {singular_name} object.
+        selector: dict or None (default None)
+            Dict to use as selection criteria.
+            {singular_name} objects will be selected if they contain
+            properties corresponding to all of the dictionary's keys, with
+            values that exactly match the supplied values. If None
+            (the default), all {singular_name} objects are selected.
+        row, col: int or None (default None)
+            Subplot row and column index of {singular_name} objects to select.
+            To select {singular_name} objects by row and column, the Figure
+            must have been created using plotly.subplots.make_subplots.
+            If None (the default), all {singular_name} objects are selected.
+        
+        Returns
+        -------
+        self
+            Returns the Figure object that the method was called on
+        \"\"\"
+        for obj in self.select_{plural_name}(
+                selector=selector, row=row, col=col):
+            fn(obj)
+
+        return self
+
+    def update_{plural_name}(self, patch, selector=None, row=None, col=None):
+        \"\"\"
+        Perform a property update operation on all {singular_name} objects
+        that satisfy the specified selection criteria
+        
+        Parameters
+        ----------
+        patch: dict
+            Dictionary of property updates to be applied to all
+            {singular_name} objects that satisfy the selection criteria.
+        selector: dict or None (default None)
+            Dict to use as selection criteria.
+            {singular_name} objects will be selected if they contain
+            properties corresponding to all of the dictionary's keys, with
+            values that exactly match the supplied values. If None
+            (the default), all {singular_name} objects are selected.
+        row, col: int or None (default None)
+            Subplot row and column index of {singular_name} objects to select.
+            To select {singular_name} objects by row and column, the Figure
+            must have been created using plotly.subplots.make_subplots.
+            If None (the default), all {singular_name} objects are selected.
+        
+        Returns
+        -------
+        self
+            Returns the Figure object that the method was called on
+        \"\"\"
+        for obj in self.select_{plural_name}(
+                selector=selector, row=row, col=col):
+            obj.update(patch)
+
+        return self""")
+
     # Return source string
     # --------------------
     buffer.write('\n')
@@ -150,7 +260,8 @@ def add_{trace_node.plotly_name}(self""")
 def write_figure_classes(outdir, trace_node,
                          data_validator,
                          layout_validator,
-                         frame_validator):
+                         frame_validator,
+                         subplot_nodes):
     """
     Construct source code for the Figure and FigureWidget classes and
     write to graph_objs/_figure.py and graph_objs/_figurewidget.py
@@ -169,6 +280,8 @@ def write_figure_classes(outdir, trace_node,
         LayoutValidator instance
     frame_validator : CompoundArrayValidator
         FrameValidator instance
+    subplot_nodes: list of str
+        List of names of all of the layout subplot properties
 
     Returns
     -------
@@ -195,7 +308,8 @@ def write_figure_classes(outdir, trace_node,
                                         fig_classname,
                                         data_validator,
                                         layout_validator,
-                                        frame_validator)
+                                        frame_validator,
+                                        subplot_nodes)
 
         # ### Format and write to file###
         filepath = opath.join(outdir, 'graph_objs',

optional-requirements.txt

@@ -23,6 +23,7 @@ psutil
 
 ## codegen dependencies ##
 yapf
+inflect
 
 ## template generation ##
 colorcet