Cast the returned iterators to generators

Author: nicholas-esterer

packages/python/plotly/plotly/basedatatypes.py

@@ -382,6 +382,12 @@ def _axis_spanning_shapes_docstr(shape_type):
     return docstr
 
 
+def _generator(i):
+    """ "cast" an iterator to a generator """
+    while True:
+        yield next(i)
+
+
 class BaseFigure(object):
     """
     Base class for all figure types (both widget and non-widget)
@@ -1186,7 +1192,7 @@ def _filter_by_subplot_ref(trace):
         if filter_by_subplot:
             funcs.append(_filter_by_subplot_ref)
 
-        return self._filter_by_selector(self.data, funcs, selector)
+        return _generator(self._filter_by_selector(self.data, funcs, selector))
 
     @staticmethod
     def _selector_matches(obj, selector):
@@ -1439,7 +1445,7 @@ def _select_layout_subplots_by_prefix(
             _natural_sort_strings(list(self.layout)),
         )
         layout_objs = [self.layout[k] for k in layout_keys]
-        return self._filter_by_selector(layout_objs, [], selector)
+        return _generator(self._filter_by_selector(layout_objs, [], selector))
 
     def _select_annotations_like(
         self, prop, selector=None, row=None, col=None, secondary_y=None
@@ -1486,7 +1492,7 @@ def _filter_sec_y(obj):
 
         funcs = [_filter_row, _filter_col, _filter_sec_y]
 
-        return self._filter_by_selector(self.layout[prop], funcs, selector)
+        return _generator(self._filter_by_selector(self.layout[prop], funcs, selector))
 
     def _add_annotation_like(
         self,

packages/python/plotly/plotly/tests/test_core/test_update_objects/test_selector_matches.py

@@ -89,3 +89,8 @@ def _sel(d):
         return d["x"] == 1 and d["y"] == 3 and d["text"] == "pat metheny"
 
     assert BaseFigure._selector_matches(obj, _sel) == False
+
+
+def test_string_selector_matches_type_key():
+    assert BaseFigure._selector_matches(dict(type="bar"), "bar")
+    assert BaseFigure._selector_matches(dict(type="scatter"), "bar") == False

packages/python/plotly/plotly/tests/test_core/test_update_objects/test_update_traces.py

@@ -260,7 +260,7 @@ def f(t):
 
     def test_select_traces_type_error(self):
         with self.assertRaises(TypeError):
-            self.assert_select_traces([0], selector=123, row=1, col=1)
+            self.assert_select_traces([0], selector=123.456, row=1, col=1)
 
     def test_for_each_trace_lowercase_names(self):
         # Names are all uppercase to start