RFC: implement attr function and "magic" underscore behavior for PlotlyDict.update

plotly/graph_objs/__init__.py

@@ -12,3 +12,5 @@
 from __future__ import absolute_import
 
 from plotly.graph_objs.graph_objs import *  # this is protected with __all__
+
+from plotly.graph_objs.graph_objs_tools import attr

plotly/graph_objs/graph_objs_tools.py

@@ -1,4 +1,5 @@
 from __future__ import absolute_import
+import re
 import textwrap
 import six
 
@@ -268,3 +269,223 @@ def sort_keys(key):
     """
     is_special = key in 'rtxyz'
     return not is_special, key
+
+
+_underscore_attr_regex = re.compile(
+    "(" + "|".join(graph_reference.UNDERSCORE_ATTRS) + ")"
+)
+
+
+def _key_parts(key):
+    """
+    Split a key containing undescores into all its parts.
+
+    This function is aware of attributes that have underscores in their
+    name (e.g. ``scatter.error_x``) and does not split them incorrectly.
+
+    Also, the function always returns a list, even if there is only one item
+    in that list (e.g. `_key_parts("marker")` would return `["marker"]`)
+
+    :param (str|unicode) key: the attribute name
+    :return: (list[str|unicode]): a list with all the parts of the attribute
+    """
+    if "_" in key:
+        match = _underscore_attr_regex.search(key)
+        if match is not None:
+            if key in graph_reference.UNDERSCORE_ATTRS:
+                # we have _exactly_ one of the underscore
+                # attrs
+                return [key]
+            else:
+                # have one underscore in the UNDERSCORE_ATTR
+                # and then at least one underscore not part
+                # of the attr. Need to break out the attr
+                # and then split the other parts
+                parts = []
+                if match.start() == 0:
+                    # UNDERSCORE_ATTR is at start of key
+                    parts.append(match.group(1))
+                else:
+                    # something comes first
+                    before = key[0:match.start()-1]
+                    parts.extend(before.split("_"))
+                    parts.append(match.group(1))
+
+                # now take care of anything that might come
+                # after the underscore attr
+                if match.end() < len(key):
+                    parts.extend(key[match.end()+1:].split("_"))
+
+                return parts
+        else:  # no underscore attributes. just split on `_`
+            return key.split("_")
+
+    else:
+        return [key]
+
+
+def _underscore_magic(parts, val, obj=None, skip_dict_check=False):
+    """
+    Set a potentially "deep" attribute of `obj` specified by a list of parent
+    keys (`parts`) to `val`.
+
+    :param (list[(str|unicode)] or str|unicode) parts: The path to the
+        attribute to be set on obj. If the argument is a string, then it will
+        first be passed to `_key_parts(key)` to construct the path and then
+        this function will be called again
+    :param val: The value the attribute should have
+    :param (dict_like) obj: A dict_like object that should have the attribute
+        set. If nothing is given, then an empty dictionary is created. If
+        a subtype of `plotly.graph_obsj.PlotlyDict` is passed, then the
+        setting of the attribute (and creation of parent nodes) will be
+        validated
+    :param (bool) skip_dict_check: Optional, default is False. If True and val
+        is a dict, then this funciton will ensure that all parent nodes are
+        created in `obj`.
+    :returns (dict_like) obj: an updated version of the `obj` argument (or
+        a newly created dict if `obj` was not passed).
+
+
+    Example:
+
+    ```
+    import plotly.graph_objs as go
+    from plotly.graph_objs.graph_objs_tools import _underscore_magic
+    layout = go.Layout()
+    _underscore_magic(["xaxis", "title"], "this is my xaxis", layout)
+    _underscore_magic("yaxis_titlefont", {"size": 10, "color": "red"}, layout)
+    print(layout)
+    ```
+
+    Results in
+
+    ```
+    {'xaxis': {'title': 'this is my xaxis'},
+     'yaxis': {'titlefont': {'color': 'red', 'size': 10}}}
+    ```
+    """
+    if obj is None:
+        obj = {}
+
+    if isinstance(parts, str):
+        return _underscore_magic(_key_parts(parts), val, obj)
+
+    if isinstance(val, dict) and not skip_dict_check:
+        return _underscore_magic_dict(parts, val, obj)
+
+    if len(parts) == 1:
+        obj[parts[0]] = val
+
+    if len(parts) == 2:
+        k1, k2 = parts
+        d1 = obj.get(k1, dict())
+        d1[k2] = val
+        obj[k1] = d1
+
+    if len(parts) == 3:
+        k1, k2, k3 = parts
+        d1 = obj.get(k1, dict())
+        d2 = d1.get(k2, dict())
+        d2[k3] = val
+        d1[k2] = d2
+        obj[k1] = d1
+
+    if len(parts) == 4:
+        k1, k2, k3, k4 = parts
+        d1 = obj.get(k1, dict())
+        d2 = d1.get(k2, dict())
+        d3 = d2.get(k3, dict())
+        d3[k4] = val
+        d2[k3] = d3
+        d1[k2] = d2
+        obj[k1] = d1
+
+    if len(parts) > 4:
+        msg = (
+            "The plotly schema shouldn't have any attributes nested"
+            " beyond level 4. Check that you are setting a valid attribute"
+        )
+        raise ValueError(msg)
+
+    return obj
+
+
+def _underscore_magic_dict(parts, val, obj=None):
+    if obj is None:
+        obj = {}
+    if not isinstance(val, dict):
+        msg = "This function is only meant to be called when val is a dict"
+        raise ValueError(msg)
+
+    # make sure obj has the key all the way up to parts
+    _underscore_magic(parts, {}, obj, True)
+
+    for key, val2 in val.items():
+        _underscore_magic(parts + [key], val2, obj)
+
+    return obj
+
+
+def attr(obj=None, **kwargs):
+    """
+    Create a nested attribute using "magic underscore" behavior
+
+    :param (dict_like) obj: A dict like container on which to set the
+        attribute. This will be modified in place. If nothing is passed an
+        empty dict is constructed and then returned. If a plotly graph object
+        is passed, all attributes will be validated.
+    :kwargs: All attributes that should be set on obj
+    :returns (dict_like): A modified version of the object passed to this
+        function
+
+    Example 1:
+
+    ```
+    from plotly.graph_objs import attr, Scatter
+    my_trace = attr(Scatter(),
+        marker=attr(size=4, symbol="diamond", line_color="red"),
+        hoverlabel_bgcolor="grey"
+    )
+    ```
+
+    Returns the following:
+
+    ```
+    {'hoverlabel': {'bgcolor': 'grey'},
+     'marker': {'line': {'color': 'red'}, 'size': 4, 'symbol': 'diamond'},
+     'type': 'scatter'}
+    ```
+
+    Example 2: incorrect attribute leads to an error
+    ```
+    from plotly.graph_objs import attr, Scatter
+    my_trace = attr(Scatter(),
+        marker_mode="markers"  # incorrect, should just be mode
+    )
+    ```
+
+    Returns an error:
+
+    ```
+    PlotlyDictKeyError: 'mode' is not allowed in 'marker'
+
+    Path To Error: ['marker']['mode']
+
+    Valid attributes for 'marker' at path ['marker'] under parents ['scatter']:
+
+        ['autocolorscale', 'cauto', 'cmax', 'cmin', 'color', 'colorbar',
+        'colorscale', 'colorsrc', 'gradient', 'line', 'maxdisplayed',
+        'opacity', 'opacitysrc', 'reversescale', 'showscale', 'size',
+        'sizemin', 'sizemode', 'sizeref', 'sizesrc', 'symbol', 'symbolsrc']
+
+    Run `<marker-object>.help('attribute')` on any of the above.
+    '<marker-object>' is the object at ['marker']
+    ```
+    """
+    if obj is None:
+        obj = dict()
+
+    for k, v in kwargs.items():
+        _underscore_magic(k, v, obj)
+
+    return obj

plotly/graph_reference.py

@@ -574,6 +574,30 @@ def _get_classes():
     return classes
 
 
+def _get_underscore_attrs():
+    """
+    Return a list of all figure attributes (on traces or layouts) that have
+    underscores in them
+    """
+
+    nms = set()
+
+    def extract_keys(x):
+        if isinstance(x, dict):
+            for val in x.values():
+                if isinstance(val, dict):
+                    extract_keys(val)
+            list(map(extract_keys, x.keys()))
+        elif isinstance(x, str):
+            nms.add(x)
+        else:
+            pass
+
+    extract_keys(GRAPH_REFERENCE["layout"]["layoutAttributes"])
+    extract_keys(GRAPH_REFERENCE["traces"])
+    return list(filter(lambda x: "_" in x and x[0] != "_", nms))
+
+
 # The ordering here is important.
 GRAPH_REFERENCE = get_graph_reference()
 
@@ -592,3 +616,5 @@ def _get_classes():
 OBJECT_NAME_TO_CLASS_NAME = {class_dict['object_name']: class_name
                              for class_name, class_dict in CLASSES.items()
                              if class_dict['object_name'] is not None}
+
+UNDERSCORE_ATTRS = _get_underscore_attrs()