diff --git a/doc/api/core/composite.rst b/doc/api/core/composite.rst index e844718cd5..b75b61f858 100644 --- a/doc/api/core/composite.rst +++ b/doc/api/core/composite.rst @@ -2,30 +2,136 @@ Composite Datasets ================== The :class:`pyvista.MultiBlock` class is a composite class to hold many -data sets which can be iterated over. +data sets which can be iterated over. ``MultiBlock`` behaves mostly like +a list, but has some dictionary-like features. -You can think of MultiBlock like lists or dictionaries as we can -iterate over this data structure by index and we can also access -blocks by their string name. +List-like Features +------------------ -.. pyvista-plot:: +Create empty composite dataset - Create empty composite dataset +.. jupyter-execute:: + :hide-code: - >>> import pyvista - >>> blocks = pyvista.MultiBlock() + # must have this here as our global backend may not be static + import pyvista + pyvista.set_plot_theme('document') + pyvista.set_jupyter_backend('pythreejs') + pyvista.global_theme.window_size = [600, 400] + pyvista.global_theme.axes.show = False + pyvista.global_theme.antialiasing = True + pyvista.global_theme.show_scalar_bar = False - Add a dataset to the collection +.. jupyter-execute:: - >>> blocks.append(pyvista.Sphere()) + import pyvista as pv + from pyvista import examples + blocks = pv.MultiBlock() + blocks - Or add a named block +Add some data to the collection. - >>> blocks["cube"] = pyvista.Cube(center=(0, 0, -1)) +.. jupyter-execute:: - Plotting the MultiBlock plots all the meshes contained by it. + blocks.append(pv.Sphere()) + blocks.append(pv.Cube(center=(0, 0, -1))) - >>> blocks.plot(smooth_shading=True) +Plotting the ``MultiBlock`` plots all the meshes contained by it. + +.. jupyter-execute:: + + blocks.plot(smooth_shading=True) + +``MultiBlock`` is list-like, so individual blocks can be accessed via +indices. + +.. jupyter-execute:: + + blocks[0] # Sphere + +The length of the block can be accessed through :func:`len` + +.. jupyter-execute:: + + len(blocks) + +or through the ``n_blocks`` attribute + +.. jupyter-execute:: + + blocks.n_blocks + +More specifically, ``MultiBlock`` is a :class:`collections.abc.MutableSequence` +and supports operations such as append, pop, insert, etc. Some of these operations +allow optional names to be provided for the dictionary like usage. + +.. jupyter-execute:: + + blocks.append(pv.Cone(), name="cone") + cone = blocks.pop(-1) # Pops Cone + blocks.reverse() + +``MultiBlock`` also supports slicing for getting or setting blocks. + +.. jupyter-execute:: + + blocks[0:2] # The Sphere and Cube objects in a new ``MultiBlock`` + + +Dictionary-like Features +------------------------ + + +``MultiBlock`` also has some dictionary features. We can set the name +of the blocks, and then access them + +.. jupyter-execute:: + + blocks = pv.MultiBlock([pv.Sphere(), pv.Cube()]) + blocks.set_block_name(0, "sphere") + blocks.set_block_name(1, "cube") + blocks["sphere"] # Sphere + +It is important to note that ``MultiBlock`` is not a dictionary and does +not enforce unique keys. Keys can also be ``None``. Extra care must be +taken to avoid problems using the dictionary-like features. + +PyVista tries to keep the keys ordered correctly when doing list operations. + +.. jupyter-execute:: + + blocks.reverse() + blocks.keys() + +The dictionary like features are useful when reading in data from a file. The +keys are often more understandable to access the data than the index. +:func:`pyvista.examples.download_cavity` is an OpenFoam dataset with a nested +``MultiBlock`` structure. There are two entries in the top-level object + +.. jupyter-execute:: + + data = examples.download_cavity() + data.keys() + +``"internalMesh"`` is a :class:`pyvista.UnstructuredGrid`. + +.. jupyter-execute:: + + data["internalMesh"] + +``"boundary"`` is another :class:`pyvista.MultiBlock`. + +.. jupyter-execute:: + + data["boundary"] + +Using the dictionary like features of :class:`pyvista.MultiBlock` allow for easier +inspection and use of the data coming from an outside source. The names of each key +correspond to human understable portions of the dataset. + +.. jupyter-execute:: + + data["boundary"].keys() Examples using this class: diff --git a/doc/conf.py b/doc/conf.py index 8ebdd60888..11bea99a5c 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -201,12 +201,15 @@ # wraps r'\.Plotter\.enable_depth_peeling$', r'\.add_scalar_bar$', - # pending refactor - r'\.MultiBlock\.next$', # called from inherited r'\.Table\.copy_meta_from$', # Type alias r'\.color_like$', + # Mixin methods from collections.abc + r'\.MultiBlock\.clear$', + r'\.MultiBlock\.count$', + r'\.MultiBlock\.index$', + r'\.MultiBlock\.remove$', } diff --git a/pyvista/core/composite.py b/pyvista/core/composite.py index b6dc30217a..5b4f129f52 100644 --- a/pyvista/core/composite.py +++ b/pyvista/core/composite.py @@ -4,9 +4,10 @@ to VTK algorithms and PyVista filtering/plotting routines. """ import collections.abc +from itertools import zip_longest import logging import pathlib -from typing import Any, List, Optional, Tuple, Union, cast +from typing import Any, Iterable, List, Optional, Tuple, Union, cast, overload import numpy as np @@ -21,22 +22,33 @@ log.setLevel('CRITICAL') -class MultiBlock(_vtk.vtkMultiBlockDataSet, CompositeFilters, DataObject): +_TypeMultiBlockLeaf = Union['MultiBlock', DataSet] + + +class MultiBlock( + _vtk.vtkMultiBlockDataSet, CompositeFilters, DataObject, collections.abc.MutableSequence +): """A composite class to hold many data sets which can be iterated over. - This wraps/extends the ``vtkMultiBlockDataSet`` class in VTK so - that we can easily plot these data sets and use the composite in a + This wraps/extends the `vtkMultiBlockDataSet + `_ class + so that we can easily plot these data sets and use the composite in a Pythonic manner. - You can think of ``MultiBlock`` like lists or dictionaries as we - can iterate over this data structure by index and we can also - access blocks by their string name. + You can think of ``MultiBlock`` like a list as we + can iterate over this data structure by index. It has some dictionary + features as we can also access blocks by their string name. + + .. versionchanged:: 0.36.0 + ``MultiBlock`` adheres more closely to being list like, and inherits + from :class:`collections.abc.MutableSequence`. Multiple nonconforming + behaviors were removed or modified. Examples -------- >>> import pyvista as pv - Create empty composite dataset + Create an empty composite dataset. >>> blocks = pv.MultiBlock() @@ -62,7 +74,7 @@ class MultiBlock(_vtk.vtkMultiBlockDataSet, CompositeFilters, DataObject): >>> blocks = pv.MultiBlock(data) >>> blocks.plot() - Iterate over the collection + Iterate over the collection. >>> for name in blocks.keys(): ... block = blocks[name] @@ -99,10 +111,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') @@ -283,7 +293,15 @@ def get_index_by_name(self, name: str) -> int: return i raise KeyError(f'Block name ({name}) not found') - def __getitem__(self, index: Union[int, str]) -> Optional['MultiBlock']: + @overload + def __getitem__(self, index: Union[int, str]) -> Optional[_TypeMultiBlockLeaf]: # noqa: D105 + ... # pragma: no cover + + @overload + def __getitem__(self, index: slice) -> 'MultiBlock': # noqa: D105 + ... # pragma: no cover + + def __getitem__(self, index): """Get a block by its index or name. If the name is non-unique then returns the first occurrence. @@ -292,21 +310,16 @@ 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 @@ -314,49 +327,112 @@ def __getitem__(self, index: Union[int, str]) -> Optional['MultiBlock']: data = wrap(data) return data - def append(self, dataset: DataSet): + def append(self, dataset: Optional[_TypeMultiBlockLeaf], name: Optional[str] = None): """Add a data set to the next block index. Parameters ---------- - dataset : pyvista.DataSet + dataset : pyvista.DataSet or pyvista.MultiBlock 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 as 'Block-{i:02}'. + 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']: - """Get a block by its index or name. + def extend(self, datasets: Iterable[_TypeMultiBlockLeaf]) -> None: + """Extend MultiBlock with an Iterable. + + If another MultiBlock object is supplied, the key names will + be preserved. + + Parameters + ---------- + datasets : Iterable[pyvista.DataSet or pyvista.MultiBlock] + Datasets to extend. + + 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_uniform = pv.MultiBlock({"uniform": examples.load_uniform()}) + >>> blocks.extend(blocks_uniform) + >>> len(blocks) + 3 + >>> blocks.keys() + ['cube', 'sphere', 'uniform'] + + """ + # Code based on collections.abc + if isinstance(datasets, MultiBlock): + for key, data in zip(datasets.keys(), datasets): + self.append(data, key) + else: + for v in datasets: + self.append(v) + + def get( + self, index: str, default: Optional[_TypeMultiBlockLeaf] = None + ) -> Optional[_TypeMultiBlockLeaf]: + """Get a block by its 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 or pyvista.MultiBlock, optional + Default to return if index is not in the multiblock. + Returns ------- - pyvista.DataSet - Dataset from the given index. + pyvista.DataSet or pyvista.MultiBlock 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 @@ -378,6 +454,7 @@ def set_block_name(self, index: int, name: str): ['cube', 'sphere', 'cone'] """ + index = range(self.n_blocks)[index] if name is None: return self.GetMetaData(index).Set(_vtk.vtkCompositeDataSet.NAME(), name) @@ -405,6 +482,7 @@ def get_block_name(self, index: int) -> Optional[str]: 'cube' """ + index = range(self.n_blocks)[index] meta = self.GetMetaData(index) if meta is not None: return meta.Get(_vtk.vtkCompositeDataSet.NAME()) @@ -432,7 +510,50 @@ 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 replace(self, index: int, dataset: Optional[_TypeMultiBlockLeaf]) -> None: + """Replace dataset at index while preserving key name. + + Parameters + ---------- + index : int + Index of the block to replace. + dataset : pyvista.DataSet or pyvista.MultiBlock + Dataset for replacing the one at index. + + Examples + -------- + >>> import pyvista as pv + >>> import numpy as np + >>> data = {"cube": pv.Cube(), "sphere": pv.Sphere(center=(2, 2, 0))} + >>> blocks = pv.MultiBlock(data) + >>> blocks.replace(1, pv.Sphere(center=(10, 10, 10))) + >>> blocks.keys() + ['cube', 'sphere'] + >>> np.allclose(blocks[1].center, [10., 10., 10.]) + True + + """ + name = self.get_block_name(index) + self[index] = dataset + self.set_block_name(index, name) + + @overload + def __setitem__( + self, index: Union[int, str], data: Optional[_TypeMultiBlockLeaf] + ): # noqa: D105 + ... # pragma: no cover + + @overload + def __setitem__( + self, index: slice, data: Iterable[Optional[_TypeMultiBlockLeaf]] + ): # noqa: D105 + ... # pragma: no cover + + def __setitem__( + self, + index, + data, + ): """Set a block with a VTK data object. To set the name simultaneously, pass a string name as the 2nd index. @@ -441,8 +562,11 @@ 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.keys() + ['Block-00', 'poly'] >>> multi['bar'] = pyvista.PolyData() >>> multi.n_blocks 3 @@ -450,41 +574,60 @@ def __setitem__(self, index: Union[Tuple[int, Optional[str]], int, str], data: D """ 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 + self.append(data, index) + return name = index + elif isinstance(index, slice): + index_iter = range(self.n_blocks)[index] + for i, (idx, d) in enumerate(zip_longest(index_iter, data)): + if idx is None: + self.insert( + index_iter[-1] + 1 + (i - len(index_iter)), d + ) # insert after last entry, increasing + elif d is None: + del self[index_iter[-1] + 1] # delete next entry + else: + self[idx] = d # + return else: - i, name = cast(int, index), None + i = index + + # data, i, and name are a single value now 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 + i = range(self.n_blocks)[i] - # check if we are overwriting a block - existing_dataset = self.GetBlock(i) - if existing_dataset is not None: - self._remove_ref(i) + # this is the only spot in the class where we actually add + # data to the MultiBlock - self.SetBlock(i, data) - if data is not None: - self._refs[data.memory_address] = data + # 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 if name is None: name = f'Block-{i:02}' self.set_block_name(i, name) # Note that this calls self.Modified() - def __delitem__(self, index: Union[int, str]): + def __delitem__(self, index: Union[int, str, slice]) -> None: """Remove a block at the specified index.""" + if isinstance(index, slice): + if index.indices(self.n_blocks)[2] > 0: + for i in reversed(range(*index.indices(self.n_blocks))): + self.__delitem__(i) + else: + for i in range(*index.indices(self.n_blocks)): + self.__delitem__(i) + return if isinstance(index, str): index = self.get_index_by_name(index) self._remove_ref(index) @@ -520,35 +663,110 @@ def __eq__(self, other): return True - def next(self) -> Optional['MultiBlock']: + def __next__(self) -> Optional[_TypeMultiBlockLeaf]: """Get the next block from the iterator.""" if self._iter_n < self.n_blocks: result = self[self._iter_n] self._iter_n += 1 return result - else: - raise StopIteration + raise StopIteration + + def insert(self, index: int, dataset: _TypeMultiBlockLeaf, name: Optional[str] = None) -> None: + """Insert data before index. + + Parameters + ---------- + index : int + Index before which to insert data. + dataset : pyvista.DataSet or pyvista.MultiBlock + Data to insert. + name : str, optional + Name for key to give dataset. A default name is given + depending on the block index as ``'Block-{i:02}'``. + + Examples + -------- + Insert a new :class:`pyvista.PolyData` at the start of the multiblock. + + >>> import pyvista as pv + >>> data = {"cube": pv.Cube(), "sphere": pv.Sphere(center=(2, 2, 0))} + >>> blocks = pv.MultiBlock(data) + >>> blocks.keys() + ['cube', 'sphere'] + >>> blocks.insert(0, pv.Plane(), "plane") + >>> blocks.keys() + ['plane', 'cube', 'sphere'] + + """ + index = range(self.n_blocks)[index] - __next__ = next + self.n_blocks += 1 + for i in reversed(range(index, self.n_blocks - 1)): + self[i + 1] = self[i] + self.set_block_name(i + 1, self.get_block_name(i)) - def pop(self, index: Union[int, str]) -> Optional['MultiBlock']: + self[index] = dataset + self.set_block_name(index, name) + + def pop(self, index: Union[int, str] = -1) -> Optional[_TypeMultiBlockLeaf]: """Pop off a block at the specified index. Parameters ---------- - index : int or str - Index or name of the dataset within the multiblock. + index : int or str, optional + Index or name of the dataset within the multiblock. Defaults to + last dataset. Returns ------- - pyvista.DataSet - Dataset from the given index. + pyvista.DataSet or pyvista.MultiBlock + Dataset from the given index that was removed. + + Examples + -------- + Pop the ``"cube"`` multiblock. + + >>> import pyvista as pv + >>> data = {"cube": pv.Cube(), "sphere": pv.Sphere(center=(2, 2, 0))} + >>> blocks = pv.MultiBlock(data) + >>> blocks.keys() + ['cube', 'sphere'] + >>> cube = blocks.pop("cube") + >>> blocks.keys() + ['sphere'] """ + if isinstance(index, int): + index = range(self.n_blocks)[index] data = self[index] del self[index] return data + def reverse(self): + """Reverse MultiBlock in-place. + + Examples + -------- + Reverse a multiblock. + + >>> import pyvista as pv + >>> data = {"cube": pv.Cube(), "sphere": pv.Sphere(center=(2, 2, 0))} + >>> blocks = pv.MultiBlock(data) + >>> blocks.keys() + ['cube', 'sphere'] + >>> blocks.reverse() + >>> blocks.keys() + ['sphere', 'cube'] + + """ + # Taken from implementation in collections.abc.MutableSequence + names = self.keys() + n = len(self) + for i in range(n // 2): + self[i], self[n - i - 1] = self[n - i - 1], self[i] + for i, name in enumerate(reversed(names)): + self.set_block_name(i, name) + def clean(self, empty=True): """Remove any null blocks in place. diff --git a/pyvista/core/filters/data_set.py b/pyvista/core/filters/data_set.py index fe4feb20ff..eccb5f4bd3 100644 --- a/pyvista/core/filters/data_set.py +++ b/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): diff --git a/pyvista/utilities/fileio.py b/pyvista/utilities/fileio.py index 9f32219549..69926ff35b 100644 --- a/pyvista/utilities/fileio.py +++ b/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): diff --git a/tests/plotting/test_plotting.py b/tests/plotting/test_plotting.py index d27fd79b68..091b1570b8 100644 --- a/tests/plotting/test_plotting.py +++ b/tests/plotting/test_plotting.py @@ -1179,7 +1179,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) diff --git a/tests/test_composite.py b/tests/test_composite.py index 8663a06722..14262aa48b 100644 --- a/tests/test_composite.py +++ b/tests/test_composite.py @@ -123,7 +123,8 @@ def test_multi_block_set_get_ers(): assert multi.n_blocks == 6 # Check pyvista side registered it # Add data to the MultiBlock data = ex.load_rectilinear() - multi[1, 'rect'] = data + multi[1] = data + multi.set_block_name(1, 'rect') # Make sure number of blocks is constant assert multi.n_blocks == 6 # Check content @@ -139,6 +140,9 @@ def test_multi_block_set_get_ers(): # Test get by name assert isinstance(multi['uni'], UniformGrid) assert isinstance(multi['rect'], RectilinearGrid) + assert isinstance(multi.get('uni'), UniformGrid) + assert multi.get('no key') is None + assert multi.get('no key', default=pyvista.Sphere()) == pyvista.Sphere() # Test the del operator del multi[0] assert multi.n_blocks == 5 @@ -153,21 +157,140 @@ def test_multi_block_set_get_ers(): pop = multi.pop(0) assert isinstance(pop, RectilinearGrid) assert multi.n_blocks == 3 - assert multi.get_block_name(10) is None + assert all([k is None for k in multi.keys()]) + + multi["new key"] = pyvista.Sphere() + assert multi.n_blocks == 4 + assert multi[3] == pyvista.Sphere() + + multi["new key"] = pyvista.Cube() + assert multi.n_blocks == 4 + assert multi[3] == pyvista.Cube() + with pytest.raises(KeyError): _ = multi.get_index_by_name('foo') - # allow Sequence but not Iterable in setitem + + with pytest.raises(IndexError): + multi[4] = UniformGrid() + + with pytest.raises(KeyError): + multi["not a key"] + with pytest.raises(TypeError): + data = multi[[0, 1]] + with pytest.raises(TypeError): - multi[{1, 'foo'}] = data + multi[1, 'foo'] = data + + +def test_replace(): + spheres = {f"{i}": pyvista.Sphere(phi_resolution=i + 3) for i in range(10)} + multi = MultiBlock(spheres) + cube = pyvista.Cube() + multi.replace(3, cube) + assert multi.get_block_name(3) == "3" + assert multi[3] is cube + + +def test_pop(): + spheres = {f"{i}": pyvista.Sphere(phi_resolution=i + 3) for i in range(10)} + multi = MultiBlock(spheres) + assert multi.pop() == spheres["9"] + assert spheres["9"] not in multi + assert multi.pop(0) == spheres["0"] + assert spheres["0"] not in multi + + +def test_del_slice(sphere): + multi = MultiBlock({f"{i}": sphere for i in range(10)}) + del multi[0:10:2] + assert len(multi) == 5 + assert all([f"{i}" in multi.keys() for i in range(1, 10, 2)]) + + multi = MultiBlock({f"{i}": sphere for i in range(10)}) + del multi[5:2:-1] + assert len(multi) == 7 + assert all([f"{i}" in multi.keys() for i in [0, 1, 2, 6, 7, 8, 9]]) + + +def test_slicing_multiple_in_setitem(sphere): + # equal length + multi = MultiBlock({f"{i}": sphere for i in range(10)}) + multi[1:3] = [pyvista.Cube(), pyvista.Cube()] + assert multi[1] == pyvista.Cube() + assert multi[2] == pyvista.Cube() + assert multi.count(pyvista.Cube()) == 2 + assert len(multi) == 10 + + # len(slice) < len(data) + multi = MultiBlock({f"{i}": sphere for i in range(10)}) + multi[1:3] = [pyvista.Cube(), pyvista.Cube(), pyvista.Cube()] + assert multi[1] == pyvista.Cube() + assert multi[2] == pyvista.Cube() + assert multi[3] == pyvista.Cube() + assert multi.count(pyvista.Cube()) == 3 + assert len(multi) == 11 + + # len(slice) > len(data) + multi = MultiBlock({f"{i}": sphere for i in range(10)}) + multi[1:3] = [pyvista.Cube()] + assert multi[1] == pyvista.Cube() + assert multi.count(pyvista.Cube()) == 1 + assert len(multi) == 9 + + +def test_reverse(sphere): + multi = MultiBlock({f"{i}": sphere for i in range(3)}) + multi.append(pyvista.Cube(), "cube") + multi.reverse() + assert multi[0] == pyvista.Cube() + assert np.array_equal(multi.keys(), ["cube", "2", "1", "0"]) + + +def test_insert(sphere): + multi = MultiBlock({f"{i}": sphere for i in range(3)}) + cube = pyvista.Cube() + multi.insert(0, cube) + assert len(multi) == 4 + assert multi[0] is cube + + # test with negative index and name + multi.insert(-1, pyvista.UniformGrid(), name="uni") + assert len(multi) == 5 + # inserted before last element + assert isinstance(multi[-2], pyvista.UniformGrid) # inserted before last element + assert multi.get_block_name(-2) == "uni" + + +def test_extend(sphere, uniform, ant): + # test with Iterable + multi = MultiBlock([sphere, ant]) + new_multi = [uniform, uniform] + multi.extend(new_multi) + assert len(multi) == 4 + assert multi.count(uniform) == 2 + + # test with a MultiBlock + multi = MultiBlock([sphere, ant]) + new_multi = MultiBlock({"uniform1": uniform, "uniform2": uniform}) + multi.extend(new_multi) + assert len(multi) == 4 + assert multi.count(uniform) == 2 + assert multi.keys()[-2] == "uniform1" + assert multi.keys()[-1] == "uniform2" def test_multi_block_clean(rectilinear, uniform, ant): # now test a clean of the null values multi = MultiBlock() - multi[1, 'rect'] = rectilinear - multi[2, 'empty'] = PolyData() - multi[3, 'mempty'] = MultiBlock() - multi[5, 'uni'] = uniform + multi.n_blocks = 6 + multi[1] = rectilinear + multi.set_block_name(1, 'rect') + multi[2] = PolyData() + multi.set_block_name(2, 'empty') + multi[3] = MultiBlock() + multi.set_block_name(3, 'mempty') + multi[5] = uniform + multi.set_block_name(5, 'uni') # perform the clean to remove all Null elements multi.clean() assert multi.n_blocks == 2 @@ -178,11 +301,15 @@ def test_multi_block_clean(rectilinear, uniform, ant): assert multi.get_block_name(1) == 'uni' # Test a nested data struct foo = MultiBlock() + foo.n_blocks = 4 foo[3] = ant assert foo.n_blocks == 4 multi = MultiBlock() - multi[1, 'rect'] = rectilinear - multi[5, 'multi'] = foo + multi.n_blocks = 6 + multi[1] = rectilinear + multi.set_block_name(1, 'rect') + multi[5] = foo + multi.set_block_name(5, 'multi') # perform the clean to remove all Null elements assert multi.n_blocks == 6 multi.clean() @@ -344,6 +471,14 @@ def test_multi_block_negative_index(ant, sphere, uniform, airplane, globe): with pytest.raises(IndexError): _ = multi[-6] + multi[-1] = ant + assert multi[4] == ant + multi[-5] = globe + assert multi[0] == globe + + with pytest.raises(IndexError): + multi[-6] = uniform + def test_multi_slice_index(ant, sphere, uniform, airplane, globe): multi = multi_from_datasets(ant, sphere, uniform, airplane, globe) @@ -365,6 +500,11 @@ def test_multi_slice_index(ant, sphere, uniform, airplane, globe): assert sub[i] is multi[j] assert sub.get_block_name(i) == multi.get_block_name(j) + sub = [airplane, globe] + multi[0:2] = sub + assert multi[0] is airplane + assert multi[1] is globe + def test_slice_defaults(ant, sphere, uniform, airplane, globe): multi = multi_from_datasets(ant, sphere, uniform, airplane, globe) @@ -386,26 +526,6 @@ def test_slice_negatives(ant, sphere, uniform, airplane, globe): assert multi[-1:-4:-2] == test_multi -def test_multi_block_list_index(ant, sphere, uniform, airplane, globe): - multi = multi_from_datasets(ant, sphere, uniform, airplane, globe) - # Now check everything - indices = [0, 3, 4] - sub = multi[indices] - assert len(sub) == len(indices) - for i, j in enumerate(indices): - assert id(sub[i]) == id(multi[j]) - assert sub.get_block_name(i) == multi.get_block_name(j) - # check list of key names - multi = MultiBlock() - multi["foo"] = pyvista.Sphere() - multi["goo"] = pyvista.Box() - multi["soo"] = pyvista.Cone() - indices = ["goo", "foo"] - sub = multi[indices] - assert len(sub) == len(indices) - assert isinstance(sub["foo"], PolyData) - - def test_multi_block_volume(ant, airplane, sphere, uniform): multi = multi_from_datasets(ant, sphere, uniform, airplane, None) vols = ant.volume + sphere.volume + uniform.volume + airplane.volume