Make MultiBlock meshes MutableSequence with some dict-like features for setting and getting.

pyvista/core/composite.py

@@ -6,7 +6,7 @@
 import collections.abc
 import logging
 import pathlib
-from typing import Any, List, Optional, Tuple, Union, cast
+from typing import Any, List, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 
@@ -99,10 +99,8 @@ def __init__(self, *args, **kwargs) -> None:
             elif isinstance(args[0], (str, pathlib.Path)):
                 self._from_file(args[0], **kwargs)
             elif isinstance(args[0], dict):
-                idx = 0
                 for key, block in args[0].items():
-                    self[idx, key] = block
-                    idx += 1
+                    self.append(block, key)
             else:
                 raise TypeError(f'Type {type(args[0])} is not supported by pyvista.MultiBlock')
 
@@ -292,71 +290,94 @@ def __getitem__(self, index: Union[int, str]) -> Optional['MultiBlock']:
         if isinstance(index, slice):
             multi = MultiBlock()
             for i in range(self.n_blocks)[index]:
-                multi[-1, self.get_block_name(i)] = self[i]
-            return multi
-        elif isinstance(index, (list, tuple, np.ndarray)):
-            multi = MultiBlock()
-            for i in index:
-                name = i if isinstance(i, str) else self.get_block_name(i)
-                multi[-1, name] = self[i]  # type: ignore
+                multi.append(self[i], self.get_block_name(i))
             return multi
         elif isinstance(index, str):
             index = self.get_index_by_name(index)
         ############################
+        if index < -self.n_blocks or index >= self.n_blocks:
+            raise IndexError(f'index ({index}) out of range for this dataset.')
         if index < 0:
             index = self.n_blocks + index
-        if index < 0 or index >= self.n_blocks:
-            raise IndexError(f'index ({index}) out of range for this dataset.')
+
         data = self.GetBlock(index)
         if data is None:
             return data
         if data is not None and not is_pyvista_dataset(data):
             data = wrap(data)
         return data
 
-    def append(self, dataset: DataSet):
+    def append(self, dataset: DataSet, name: Optional[str] = None):
         """Add a data set to the next block index.
 
         Parameters
         ----------
         dataset : pyvista.DataSet
             Dataset to append to this multi-block.
 
+        name : str, optional
+            Block name to give to dataset.  A default name is given
+            depending on the block index.
+
         Examples
         --------
         >>> import pyvista as pv
+        >>> from pyvista import examples
         >>> data = {"cube": pv.Cube(), "sphere": pv.Sphere(center=(2, 2, 0))}
         >>> blocks = pv.MultiBlock(data)
         >>> blocks.append(pv.Cone())
         >>> len(blocks)
         3
+        >>> blocks.append(examples.load_uniform(), "uniform")
+        >>> blocks.keys()
+        ['cube', 'sphere', 'Block-02', 'uniform']
 
         """
         index = self.n_blocks  # note off by one so use as index
         # always wrap since we may need to reference the VTK memory address
         if not pyvista.is_pyvista_dataset(dataset):
             dataset = pyvista.wrap(dataset)
+        self.n_blocks += 1
         self[index] = dataset
+        # No overwrite if name is None
+        self.set_block_name(index, name)
 
-    def get(self, index: Union[int, str]) -> Optional['MultiBlock']:
+    def get(self, index: str, default: Optional[DataSet] = None) -> Optional[DataSet]:
         """Get a block by its index or name.
 
         If the name is non-unique then returns the first occurrence.
+        Returns ``default`` if name isn't in the dataset.
 
         Parameters
         ----------
-        index : int or str
+        index : str
             Index or name of the dataset within the multiblock.
 
+        default : pyvista.DataSet, optional
+            Default to return if index is not in the multiblock.
+
         Returns
         -------
-        pyvista.DataSet
-            Dataset from the given index.
+        pyvista.DataSet or None
+            Dataset from the given index if it exists.
+
+        Examples
+        --------
+        >>> import pyvista as pv
+        >>> from pyvista import examples
+        >>> data = {"poly": pv.PolyData(), "uni": pv.UniformGrid()}
+        >>> blocks = pv.MultiBlock(data)
+        >>> blocks.get("poly")
+        PolyData ...
+        >>> blocks.get("cone")
 
         """
-        return self[index]
+        try:
+            return self[index]
+        except KeyError:
+            return default
 
-    def set_block_name(self, index: int, name: str):
+    def set_block_name(self, index: int, name: Optional[str]):
         """Set a block's string name at the specified index.
 
         Parameters
@@ -432,7 +453,11 @@ def keys(self) -> List[Optional[str]]:
     def _ipython_key_completions_(self) -> List[Optional[str]]:
         return self.keys()
 
-    def __setitem__(self, index: Union[Tuple[int, Optional[str]], int, str], data: DataSet):
+    def __setitem__(
+        self,
+        index: Union[int, str, slice],
+        data: Union[DataSet, Sequence[DataSet], Tuple[str, DataSet]],
+    ):
         """Set a block with a VTK data object.
 
         To set the name simultaneously, pass a string name as the 2nd index.
@@ -441,43 +466,67 @@ def __setitem__(self, index: Union[Tuple[int, Optional[str]], int, str], data: D
         -------
         >>> import pyvista
         >>> multi = pyvista.MultiBlock()
-        >>> multi[0] = pyvista.PolyData()
-        >>> multi[1, 'foo'] = pyvista.UnstructuredGrid()
+        >>> multi.append(pyvista.PolyData())
+        >>> multi[0] = pyvista.UnstructuredGrid()
+        >>> multi.append(pyvista.PolyData(), 'poly')
+        >>> multi[1] = ('foo', pyvista.UnstructuredGrid())
+        >>> multi.keys()
+        ['Block-00', 'foo']
         >>> multi['bar'] = pyvista.PolyData()
         >>> multi.n_blocks
         3
 
         """
+        if isinstance(index, str) and isinstance(data, tuple):
+            raise TypeError(f"Cannot set key {index} with a different string key from {data}")
+
         i: int = 0
         name: Optional[str] = None
-        if isinstance(index, (np.ndarray, collections.abc.Sequence)) and not isinstance(index, str):
-            i, name = index[0], index[1]
-        elif isinstance(index, str):
+        if isinstance(index, str):
             try:
                 i = self.get_index_by_name(index)
             except KeyError:
-                i = -1
+                # data cannot be a tuple here
+                if isinstance(data, collections.abc.Sequence):
+                    raise TypeError(f"Cannot set key {index} with the sequence {data}")
+                self.append(data, index)
+                return
             name = index
+        elif isinstance(index, slice):
+            if isinstance(data, tuple):
+                raise TypeError(f"Cannot set the slice {slice} with a tuple {tuple}")
+            for i, d in zip(range(self.n_blocks)[index], data):
+                self[i] = d
+            return
         else:
-            i, name = cast(int, index), None
+            i = index
+
+        # data, i, and name are a single value now
+        if isinstance(data, tuple):
+            if not len(data) == 2:
+                raise ValueError(f"Data {data} must be a length 2 tuple of form (DataSet, str)")
+            name, data = data
+
         if data is not None and not is_pyvista_dataset(data):
             data = wrap(data)
+        data = cast(pyvista.DataSet, data)
 
-        if i == -1:
-            self.append(data)
-            i = self.n_blocks - 1
-        else:
-            # this is the only spot in the class where we actually add
-            # data to the MultiBlock
+        if i < -self.n_blocks or i >= self.n_blocks:
+            raise IndexError(f'index ({i}) out of range for this dataset.')
+        if i < 0:
+            i = self.n_blocks + i
+
+        # this is the only spot in the class where we actually add
+        # data to the MultiBlock
 
-            # check if we are overwriting a block
-            existing_dataset = self.GetBlock(i)
-            if existing_dataset is not None:
-                self._remove_ref(i)
+        # check if we are overwriting a block
+        existing_dataset = self.GetBlock(i)
+        if existing_dataset is not None:
+            self._remove_ref(i)
 
-            self.SetBlock(i, data)
-            if data is not None:
-                self._refs[data.memory_address] = data
+        self.SetBlock(i, data)
+        if data is not None:
+            self._refs[data.memory_address] = data
 
         if name is None:
             name = f'Block-{i:02}'

pyvista/core/filters/data_set.py

@@ -641,27 +641,38 @@ def slice_orthogonal(
         output = pyvista.MultiBlock()
         if isinstance(self, pyvista.MultiBlock):
             for i in range(self.n_blocks):
-                output[i] = self[i].slice_orthogonal(
-                    x=x, y=y, z=z, generate_triangles=generate_triangles, contour=contour
+                output.append(
+                    self[i].slice_orthogonal(
+                        x=x, y=y, z=z, generate_triangles=generate_triangles, contour=contour
+                    )
                 )
             return output
-        output[0, 'YZ'] = self.slice(
-            normal='x',
-            origin=[x, y, z],
-            generate_triangles=generate_triangles,
-            progress_bar=progress_bar,
+        output.append(
+            self.slice(
+                normal='x',
+                origin=[x, y, z],
+                generate_triangles=generate_triangles,
+                progress_bar=progress_bar,
+            ),
+            'YZ',
         )
-        output[1, 'XZ'] = self.slice(
-            normal='y',
-            origin=[x, y, z],
-            generate_triangles=generate_triangles,
-            progress_bar=progress_bar,
+        output.append(
+            self.slice(
+                normal='y',
+                origin=[x, y, z],
+                generate_triangles=generate_triangles,
+                progress_bar=progress_bar,
+            ),
+            'XZ',
         )
-        output[2, 'XY'] = self.slice(
-            normal='z',
-            origin=[x, y, z],
-            generate_triangles=generate_triangles,
-            progress_bar=progress_bar,
+        output.append(
+            self.slice(
+                normal='z',
+                origin=[x, y, z],
+                generate_triangles=generate_triangles,
+                progress_bar=progress_bar,
+            ),
+            'XY',
         )
         return output
 
@@ -767,14 +778,16 @@ def slice_along_axis(
         output = pyvista.MultiBlock()
         if isinstance(self, pyvista.MultiBlock):
             for i in range(self.n_blocks):
-                output[i] = self[i].slice_along_axis(
-                    n=n,
-                    axis=ax_label,
-                    tolerance=tolerance,
-                    generate_triangles=generate_triangles,
-                    contour=contour,
-                    bounds=bounds,
-                    center=center,
+                output.append(
+                    self[i].slice_along_axis(
+                        n=n,
+                        axis=ax_label,
+                        tolerance=tolerance,
+                        generate_triangles=generate_triangles,
+                        contour=contour,
+                        bounds=bounds,
+                        center=center,
+                    )
                 )
             return output
         for i in range(n):
@@ -787,7 +800,7 @@ def slice_along_axis(
                 contour=contour,
                 progress_bar=progress_bar,
             )
-            output[i, f'slice{i}'] = slc
+            output.append(slc, f'slice{i}')
         return output
 
     def slice_along_line(self, line, generate_triangles=False, contour=False, progress_bar=False):

pyvista/utilities/fileio.py

@@ -152,7 +152,7 @@ def read(filename, attrs=None, force_ext=None, file_format=None, progress_bar=Fa
                 name = os.path.basename(str(each))
             else:
                 name = None
-            multi[-1, name] = read(each, attrs=attrs, file_format=file_format)
+            multi.append(read(each, attrs=attrs, file_format=file_format), name)
         return multi
     filename = os.path.abspath(os.path.expanduser(str(filename)))
     if not os.path.isfile(filename):

tests/plotting/test_plotting.py

@@ -1165,7 +1165,7 @@ def test_multi_block_plot():
     uni.cell_data.set_array(arr, 'Random Data')
     multi.append(uni)
     # And now add a data set without the desired array and a NULL component
-    multi[3] = examples.load_airplane()
+    multi.append(examples.load_airplane())
     with pytest.raises(KeyError):
         # The scalars are not available in all datasets so raises KeyError
         multi.plot(scalars='Random Data', multi_colors=True)