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

Author: sglyon

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.py

@@ -612,7 +612,12 @@ def update(self, dict1=None, **dict2):
                     else:
                         self[key] = val
                 else:
-                    self[key] = val
+                    # don't have this key -- might be using underscore magic
+                    graph_objs_tools._underscore_magic(key, val, self)
+
+        # return self so we can chain this method (e.g. Scatter().update(**)
+        # returns an instance of Scatter)
+        return self
 
     def strip_style(self):
         """

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,117 @@ 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):
+    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):
+    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):
+    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,26 @@ def _get_classes():
     return classes
 
 
+def _get_underscore_attrs():
+
+    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 +612,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()

plotly/tests/test_core/test_graph_objs/test_graph_objs_tools.py

@@ -4,6 +4,7 @@
 
 from plotly import graph_reference as gr
 from plotly.graph_objs import graph_objs_tools as got
+from plotly import graph_objs as go
 
 
 class TestGetHelp(TestCase):
@@ -30,3 +31,183 @@ def test_get_help_does_not_raise(self):
                 got.get_help(object_name, attribute=fake_attribute)
         except:
             self.fail(msg=msg)
+
+
+class TestKeyParts(TestCase):
+    def test_without_underscore_attr(self):
+        assert got._key_parts("foo") == ["foo"]
+        assert got._key_parts("foo_bar") == ["foo", "bar"]
+        assert got._key_parts("foo_bar_baz") == ["foo", "bar", "baz"]
+
+    def test_traililng_underscore_attr(self):
+        assert got._key_parts("foo_error_x") == ["foo", "error_x"]
+        assert got._key_parts("foo_bar_error_x") == ["foo", "bar", "error_x"]
+        assert got._key_parts("foo_bar_baz_error_x") == ["foo", "bar", "baz", "error_x"]
+
+    def test_leading_underscore_attr(self):
+        assert got._key_parts("error_x_foo") == ["error_x", "foo"]
+        assert got._key_parts("error_x_foo_bar") == ["error_x", "foo", "bar"]
+        assert got._key_parts("error_x_foo_bar_baz") == ["error_x", "foo", "bar", "baz"]
+
+
+class TestUnderscoreMagicDictObj(TestCase):
+
+    def test_can_split_string_key_into_parts(self):
+        obj1 = {}
+        obj2 = {}
+        got._underscore_magic("marker_line_width", 42, obj1)
+        got._underscore_magic(["marker", "line", "width"], 42, obj2)
+        want = {"marker": {"line": {"width": 42}}}
+        assert obj1 == obj2 == want
+
+    def test_will_make_tree_with_empty_dict_val(self):
+        obj = {}
+        got._underscore_magic("marker_colorbar_tickfont", {}, obj)
+        assert obj == {"marker": {"colorbar": {"tickfont": {}}}}
+
+    def test_can_set_at_depths_1to4(self):
+        # 1 level
+        obj = {}
+        got._underscore_magic("opacity", 0.9, obj)
+        assert obj == {"opacity": 0.9}
+
+        # 2 levels
+        got._underscore_magic("line_width", 10, obj)
+        assert obj == {"opacity": 0.9, "line": {"width": 10}}
+
+        # 3 levels
+        got._underscore_magic("hoverinfo_font_family", "Times", obj)
+        want = {
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}}
+        }
+        assert obj == want
+
+        # 4 levels
+        got._underscore_magic("marker_colorbar_tickfont_family", "Times", obj)
+        want = {
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}},
+            "marker": {"colorbar": {"tickfont": {"family": "Times"}}},
+        }
+        assert obj == want
+
+    def test_does_not_displace_existing_fields(self):
+        obj = {}
+        got._underscore_magic("marker_size", 10, obj)
+        got._underscore_magic("marker_line_width", 0.4, obj)
+        assert obj == {"marker": {"size": 10, "line": {"width": 0.4}}}
+
+    def test_doesnt_mess_up_underscore_attrs(self):
+        obj = {}
+        got._underscore_magic("error_x_color", "red", obj)
+        got._underscore_magic("error_x_width", 4, obj)
+        assert obj == {"error_x": {"color": "red", "width": 4}}
+
+
+class TestUnderscoreMagicPlotlyDictObj(TestCase):
+
+    def test_can_split_string_key_into_parts(self):
+        obj1 = go.Scatter()
+        obj2 = go.Scatter()
+        got._underscore_magic("marker_line_width", 42, obj1)
+        got._underscore_magic(["marker", "line", "width"], 42, obj2)
+        want = go.Scatter({"marker": {"line": {"width": 42}}})
+        assert obj1 == obj2 == want
+
+    def test_will_make_tree_with_empty_dict_val(self):
+        obj = go.Scatter()
+        got._underscore_magic("marker_colorbar_tickfont", {}, obj)
+        want = go.Scatter({"marker": {"colorbar": {"tickfont": {}}}})
+        assert obj == want
+
+    def test_can_set_at_depths_1to4(self):
+        # 1 level
+        obj = go.Scatter()
+        got._underscore_magic("opacity", 0.9, obj)
+        assert obj == go.Scatter({"type": "scatter", "opacity": 0.9})
+
+        # 2 levels
+        got._underscore_magic("line_width", 10, obj)
+        assert obj == go.Scatter({"opacity": 0.9, "line": {"width": 10}})
+
+        # 3 levels
+        got._underscore_magic("hoverinfo_font_family", "Times", obj)
+        want = go.Scatter({
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}}
+        })
+        assert obj == want
+
+        # 4 levels
+        got._underscore_magic("marker_colorbar_tickfont_family", "Times", obj)
+        want = go.Scatter({
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}},
+            "marker": {"colorbar": {"tickfont": {"family": "Times"}}},
+        })
+        assert obj == want
+
+    def test_does_not_displace_existing_fields(self):
+        obj = go.Scatter()
+        got._underscore_magic("marker_size", 10, obj)
+        got._underscore_magic("marker_line_width", 0.4, obj)
+        assert obj == go.Scatter({"marker": {"size": 10, "line": {"width": 0.4}}})
+
+    def test_doesnt_mess_up_underscore_attrs(self):
+        obj = go.Scatter()
+        got._underscore_magic("error_x_color", "red", obj)
+        got._underscore_magic("error_x_width", 4, obj)
+        assert obj == go.Scatter({"error_x": {"color": "red", "width": 4}})
+
+
+class TestAttr(TestCase):
+    def test_with_no_positional_argument(self):
+        have = got.attr(
+            opacity=0.9, line_width=10,
+            hoverinfo_font_family="Times",
+            marker_colorbar_tickfont_size=10
+        )
+        want = {
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}},
+            "marker": {"colorbar": {"tickfont": {"size": 10}}},
+        }
+        assert have == want
+
+    def test_with_dict_positional_argument(self):
+        have = {"x": [1, 2, 3, 4, 5]}
+        got.attr(have,
+            opacity=0.9, line_width=10,
+            hoverinfo_font_family="Times",
+            marker_colorbar_tickfont_size=10
+        )
+        want = {
+            "x": [1, 2, 3, 4, 5],
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}},
+            "marker": {"colorbar": {"tickfont": {"size": 10}}},
+        }
+        assert have == want
+
+    def test_with_PlotlyDict_positional_argument(self):
+        have = go.Scatter({"x": [1, 2, 3, 4, 5]})
+        got.attr(have,
+            opacity=0.9, line_width=10,
+            hoverinfo_font_family="Times",
+            marker_colorbar_tickfont_size=10
+        )
+        want = go.Scatter({
+            "x": [1, 2, 3, 4, 5],
+            "opacity": 0.9,
+            "line": {"width": 10},
+            "hoverinfo": {"font": {"family": "Times"}},
+            "marker": {"colorbar": {"tickfont": {"size": 10}}},
+        })
+        assert have == want