Fixed DynamicMap collation for changing keys

holoviews/core/spaces.py

@@ -3,6 +3,7 @@
 from numbers import Number
 from itertools import groupby
 from functools import partial
+from collections import defaultdict
 from contextlib import contextmanager
 from inspect import ArgSpec
 
@@ -1229,6 +1230,7 @@ def collate(self):
             return self
 
         container = initialized.last.clone(shared_data=False)
+        type_counter = defaultdict(int)
 
         # Get stream mapping from callback
         remapped_streams = []
@@ -1253,11 +1255,38 @@ def collate(self):
 
             # Define collation callback
             def collation_cb(*args, **kwargs):
-                return self[args][kwargs['selection_key']]
-            callback = Callable(partial(collation_cb, selection_key=k),
+                layout = self[args]
+                layout_type = type(layout).__name__
+                print(container.keys(), layout.keys())
+                if len(container.keys()) != len(layout.keys()):
+                    raise ValueError('Collated DynamicMaps must return '
+                                     '%s with consistent number of items.'
+                                     % layout_type)
+
+                key = kwargs['selection_key']
+                index = kwargs['selection_index']
+                obj_type = kwargs['selection_type']
+                dyn_type_map = defaultdict(list)
+                for k, v in layout.data.items():
+                    if k == key:
+                        return layout[k]
+                    dyn_type_map[type(v)].append(v)
+
+                dyn_type_counter = {t: len(vals) for t, vals in dyn_type_map.items()}
+                if dyn_type_counter != type_counter:
+                    raise ValueError('The objects in a %s returned by a '
+                                     'DynamicMap must consistently return '
+                                     'the same number of items of the '
+                                     'same type.' % layout_type)
+                return dyn_type_map[obj_type][index]
+
+            callback = Callable(partial(collation_cb, selection_key=k,
+                                        selection_index=type_counter[type(v)],
+                                        selection_type=type(v)),
                                 inputs=[self])
             vdmap = self.clone(callback=callback, shared_data=False,
                                streams=vstreams)
+            type_counter[type(v)] += 1
 
             # Remap source of streams
             for stream in vstreams:

tests/core/testdynamic.py

@@ -805,3 +805,55 @@ def callback():
         self.assertEqual(list(grid.keys()), [(i, j) for i in range(1, 3)
                                              for j in range(1, 3)])
         self.assertEqual(stream.source, grid[(1, 2)])
+
+    def test_dynamic_collate_layout_with_changing_label(self):
+        def callback(i):
+            return Layout([Curve([], label=str(j)) for j in range(i, i+2)])
+        dmap = DynamicMap(callback, kdims=['i']).redim.range(i=(0, 10))
+        layout = dmap.collate()
+        dmap1, dmap2 = layout.values()
+        el1, el2 = dmap1[2], dmap2[2]
+        self.assertEqual(el1.label, '2')
+        self.assertEqual(el2.label, '3')
+
+    def test_dynamic_collate_ndlayout_with_changing_keys(self):
+        def callback(i):
+            return NdLayout({j: Curve([], label=str(j)) for j in range(i, i+2)})
+        dmap = DynamicMap(callback, kdims=['i']).redim.range(i=(0, 10))
+        layout = dmap.collate()
+        dmap1, dmap2 = layout.values()
+        el1, el2 = dmap1[2], dmap2[2]
+        self.assertEqual(el1.label, '2')
+        self.assertEqual(el2.label, '3')
+
+    def test_dynamic_collate_gridspace_with_changing_keys(self):
+        def callback(i):
+            return GridSpace({j: Curve([], label=str(j)) for j in range(i, i+2)}, 'X')
+        dmap = DynamicMap(callback, kdims=['i']).redim.range(i=(0, 10))
+        layout = dmap.collate()
+        dmap1, dmap2 = layout.values()
+        el1, el2 = dmap1[2], dmap2[2]
+        self.assertEqual(el1.label, '2')
+        self.assertEqual(el2.label, '3')
+
+    def test_dynamic_collate_gridspace_with_changing_items_raises(self):
+        def callback(i):
+            return GridSpace({j: Curve([], label=str(j)) for j in range(i)}, 'X')
+        dmap = DynamicMap(callback, kdims=['i']).redim.range(i=(2, 10))
+        layout = dmap.collate()
+        dmap1, dmap2 = layout.values()
+        err = 'Collated DynamicMaps must return GridSpace with consistent number of items.'
+        with self.assertRaisesRegexp(ValueError, err):
+            dmap1[4]
+
+    def test_dynamic_collate_gridspace_with_changing_item_types_raises(self):
+        def callback(i):
+            eltype = Image if i%2 else Curve
+            return GridSpace({j: eltype([], label=str(j)) for j in range(i, i+2)}, 'X')
+        dmap = DynamicMap(callback, kdims=['i']).redim.range(i=(2, 10))
+        layout = dmap.collate()
+        dmap1, dmap2 = layout.values()
+        err = ('The objects in a GridSpace returned by a DynamicMap must '
+               'consistently return the same number of items of the same type.')
+        with self.assertRaisesRegexp(ValueError, err):
+            dmap1[3]