Repository: pyvista/pyvista

Issue:
Cannot serialize MultiBlock (dataset of type 13) with pickle
### Describe the bug, what's wrong, and what you expected.

Attempting to serialize a pyvista.MultiBlock results in a TypeError being thrown.

I expected the serialization to proceed.

Perhaps a loop needs to be added to `pyvista.core.DataObject.__getstate__`?

### Steps to reproduce the bug.

```python
import pickle
import pyvista as pv

multi_block = pv.MultiBlock((pv.Cube(), pv.Sphere()))
with open('failed.pkl', 'wb') as f:
    pickle.dump(multi_block, f)
```

```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "c:\Users\kalin\Documents\git-projects\knee-mesher\.env\Lib\site-packages\pyvista\core\dataobject.py", line 728, in __getstate__
    raise TypeError(f'Cannot pickle dataset of type {self.GetDataObjectType()}')
TypeError: Cannot pickle dataset of type 13
```

### System Information

```shell
--------------------------------------------------------------------------------
  Date: Fri Jul 26 13:30:03 2024 Mountain Daylight Time

                OS : Windows (10 10.0.19045 SP0 Multiprocessor Free)
            CPU(s) : 12
           Machine : AMD64
      Architecture : 64bit
               RAM : 31.9 GiB
       Environment : Python
        GPU Vendor : NVIDIA Corporation
      GPU Renderer : NVIDIA GeForce RTX 3070 Ti/PCIe/SSE2
       GPU Version : 4.5.0 NVIDIA 560.70
  MathText Support : False

  Python 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:27:10) [MSC
  v.1938 64 bit (AMD64)]

           pyvista : 0.44.1
               vtk : 9.3.1
             numpy : 1.26.4
        matplotlib : 3.9.1
            scooby : 0.10.0
             pooch : 1.8.2
            pillow : 10.4.0
           imageio : 2.34.2
             PyQt5 : 5.15.9
           IPython : 8.26.0
        ipywidgets : 8.1.3
             scipy : 1.14.0
              tqdm : 4.66.4
            meshio : 5.3.5
        jupyterlab : 4.2.4
      nest_asyncio : 1.6.0
--------------------------------------------------------------------------------
```


### Screenshots

_No response_

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyvista/pyvista/pull/6622?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyvista) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 97.13%. Comparing base [(`6647524`)](https://app.codecov.io/gh/pyvista/pyvista/commit/664752494d7ba82a1d79fae72c2a7e1515c95e7c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyvista) to head [(`f4a6c1a`)](https://app.codecov.io/gh/pyvista/pyvista/commit/f4a6c1aa36df45c445e66d9b5f0402ea2a374c5d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyvista).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #6622      +/-   ##
==========================================
- Coverage   97.37%   97.13%   -0.25%     
==========================================
  Files         143      143              
  Lines       27786    27786              
==========================================
- Hits        27058    26991      -67     
- Misses        728      795      +67     
```

</details>
pyvista-bot: <!-- NETLIFY DEPLOY COMMENT GENERATED BY ACTIONS_NETLIFY - APP ID SHA256: d3574f413fe51079937885070831e027ca5b10f7facb9e609494e0a481a31f79 -->
🚀 Deployed on https://66db3c5a21f96acd899da269--pyvista-dev.netlify.app
user27182: Looks like there are still issues with serializing with serializing on older vtk versions.

Patch:
diff --git a/pyvista/__init__.py b/pyvista/__init__.py
index 479a42ad..fecd55e2 100755
--- a/pyvista/__init__.py
+++ b/pyvista/__init__.py
@@ -59,7 +59,7 @@ PLOT_DIRECTIVE_THEME = None
 FLOAT_FORMAT = "{:.3e}"
 
 # Serialization format to be used when pickling `DataObject`
-PICKLE_FORMAT = 'xml'
+PICKLE_FORMAT = 'vtk' if vtk_version_info >= (9, 3) else 'xml'
 
 # Name used for unnamed scalars
 DEFAULT_SCALARS_NAME = 'Data'
diff --git a/pyvista/core/_vtk_core.py b/pyvista/core/_vtk_core.py
index ca547de0..32205387 100644
--- a/pyvista/core/_vtk_core.py
+++ b/pyvista/core/_vtk_core.py
@@ -45,6 +45,10 @@ from vtkmodules.util.numpy_support import get_vtk_array_type
 from vtkmodules.util.numpy_support import numpy_to_vtk
 from vtkmodules.util.numpy_support import numpy_to_vtkIdTypeArray
 from vtkmodules.util.numpy_support import vtk_to_numpy
+
+with contextlib.suppress(ImportError):
+    from vtkmodules.util.pickle_support import serialize_VTK_data_object
+
 from vtkmodules.util.vtkAlgorithm import VTKPythonAlgorithmBase
 from vtkmodules.vtkCommonComputationalGeometry import vtkKochanekSpline
 from vtkmodules.vtkCommonComputationalGeometry import vtkParametricBohemianDome
diff --git a/pyvista/core/composite.py b/pyvista/core/composite.py
index 6e7552b7..f11d1634 100644
--- a/pyvista/core/composite.py
+++ b/pyvista/core/composite.py
@@ -1078,7 +1078,8 @@ class MultiBlock(
         # Deep copy will not copy the block name for None blocks (name is set to None instead)
         # Iterate through the blocks to fix this recursively
         def _set_name_for_none_blocks(this_object_, new_object_):
-            for i, dataset in enumerate(pyvista.wrap(new_object_)):
+            new_object_ = pyvista.wrap(new_object_)
+            for i, dataset in enumerate(new_object_):
                 if dataset is None:
                     this_object_.set_block_name(i, new_object_.get_block_name(i))
                 elif isinstance(dataset, MultiBlock):
diff --git a/pyvista/core/dataobject.py b/pyvista/core/dataobject.py
index 16e5b27c..58d0ec2f 100644
--- a/pyvista/core/dataobject.py
+++ b/pyvista/core/dataobject.py
@@ -699,12 +699,45 @@ class DataObject:
         self.CopyAttributes(dataset)
 
     def __getstate__(self):
+        """Support pickle."""
+        pickle_format = pyvista.PICKLE_FORMAT
+        if pickle_format == 'vtk':
+            return self._serialize_vtk_pickle_format()
+        elif pickle_format in ['xml', 'legacy']:
+            return self._serialize_pyvista_pickle_format()
+        # Invalid format, use the setter to raise an error
+        pyvista.set_pickle_format(pickle_format)  # noqa:  RET503
+
+    def _serialize_vtk_pickle_format(self):
+        # Note: The serialized state has format: ( function, (dict,) )
+        serialized = _vtk.serialize_VTK_data_object(self)
+
+        # Add this object's data to the state dictionary
+        state_dict = serialized[1][0]
+        state_dict['_PYVISTA_STATE_DICT'] = self.__dict__.copy()
+
+        # Unlike the PyVista formats, we do not return a dict. Instead, return
+        # the same format returned by the vtk serializer.
+        return serialized
+
+    def _serialize_pyvista_pickle_format(self):
         """Support pickle by serializing the VTK object data to something which can be pickled natively.
 
         The format of the serialized VTK object data depends on `pyvista.PICKLE_FORMAT` (case-insensitive).
         - If `pyvista.PICKLE_FORMAT == 'xml'`, the data is serialized as an XML-formatted string.
         - If `pyvista.PICKLE_FORMAT == 'legacy'`, the data is serialized to bytes in VTK's binary format.
+
+        .. note::
+
+            These formats are custom PyVista legacy formats. The native 'vtk' format is
+            preferred since it supports more objects (e.g. MultiBlock).
+
         """
+        if isinstance(self, pyvista.MultiBlock):
+            raise TypeError(
+                "MultiBlock is not supported with 'xml' or 'legacy' pickle formats."
+                "\nUse `pyvista.PICKLE_FORMAT='vtk'`."
+            )
         state = self.__dict__.copy()
 
         if pyvista.PICKLE_FORMAT.lower() == 'xml':
@@ -751,6 +784,47 @@ class DataObject:
 
     def __setstate__(self, state):
         """Support unpickle."""
+
+        def _is_vtk_format(state_):
+            # Note: The vtk state has format ( function, (dict,) )
+            return (
+                isinstance(state_, tuple)
+                and len(state_) == 2
+                and isinstance(state_[1], tuple)
+                and len(state_[1]) == 1
+                and isinstance(state_[1][0], dict)
+            )
+
+        def _is_pyvista_format(state_):
+            return isinstance(state_, dict) and 'vtk_serialized' in state_
+
+        if _is_vtk_format(state):
+            self._unserialize_vtk_pickle_format(state)
+        elif _is_pyvista_format(state):
+            self._unserialize_pyvista_pickle_format(state)
+        else:
+            raise RuntimeError(
+                f"Cannot unpickle '{self.__class__.__name__}'. Invalid pickle format."
+            )
+
+    def _unserialize_vtk_pickle_format(self, state):
+        """Support unpickle of VTK's format."""
+        # The vtk state has format: ( function, (dict,) )
+        unserialize_func = state[0]
+        state_dict = state[1][0]
+        self.__dict__.update(state_dict['_PYVISTA_STATE_DICT'])
+        obj = unserialize_func(state_dict)
+        self.deep_copy(obj)
+
+    def _unserialize_pyvista_pickle_format(self, state):
+        """Support unpickle of PyVista 'xml' and 'legacy' formats.
+
+        .. note::
+
+            These formats are custom PyVista legacy formats. The native 'vtk' format is
+            preferred since it supports more objects (e.g. MultiBlock).
+
+        """
         vtk_serialized = state.pop('vtk_serialized')
         pickle_format = state.pop(
             'PICKLE_FORMAT',
diff --git a/pyvista/core/utilities/arrays.py b/pyvista/core/utilities/arrays.py
index 49d249b1..7aaf69d4 100644
--- a/pyvista/core/utilities/arrays.py
+++ b/pyvista/core/utilities/arrays.py
@@ -858,6 +858,22 @@ class _SerializedDictArray(UserDict, _vtk.vtkStringArray):  # type: ignore[type-
         # shows this array as `str`
         _set_string_scalar_object_name(self)
 
+    def __getstate__(self):
+        """Support pickling.
+
+        This method does nothing. It only exists to make the pickle library happy.
+        Classes that store an instance of this class must pickle this array directly.
+        E.g. DataObjects can support this by storing this array as field data
+        """
+
+    def __setstate__(self, state):
+        """Support pickling.
+
+        This method does nothing. It only exists to make the pickle library happy.
+        Classes that store an instance of this class must pickle this array directly.
+        E.g. DataObjects can support this by storing this array as field data
+        """
+
     # Override any/all `UserDict` or `MutableMapping` methods which mutate
     # the dictionary. This ensures the serialized string is also updated
     # and synced with the dict
diff --git a/pyvista/core/utilities/fileio.py b/pyvista/core/utilities/fileio.py
index c01ceffc..4a3f897f 100644
--- a/pyvista/core/utilities/fileio.py
+++ b/pyvista/core/utilities/fileio.py
@@ -21,7 +21,22 @@ def set_pickle_format(format: str):  # noqa: A002
     Parameters
     ----------
     format : str
-        The format for serialization. Acceptable values are "xml" or "legacy".
+        The format for serialization. Acceptable values are:
+
+        - ``'vtk'`` (default) : objects are serialized using VTK's official
+            marshalling methods.
+        - ``'xml'``: objects are serialized as an XML-formatted string.
+        - ``'legacy'`` objects are serialized to bytes in VTK's binary format.
+
+        .. note::
+
+            The ``'vtk'`` format requires VTK 9.3 or greater.
+
+        .. warning::
+
+            ``'xml'`` and ``'legacy'`` are not recommended. These formats are not
+            officially supported by VTK and have limitations. For example, these
+            formats cannot be used to pickle :class:`pyvista.MultiBlock`.
 
     Raises
     ------
@@ -29,12 +44,15 @@ def set_pickle_format(format: str):  # noqa: A002
         If the provided format is not supported.
 
     """
-    supported = {'xml', 'legacy'}
+    supported = {'vtk', 'xml', 'legacy'}
     format_ = format.lower()
     if format_ not in supported:
         raise ValueError(
             f'Unsupported pickle format `{format_}`. Valid options are `{"`, `".join(supported)}`.',
         )
+    if format_ == 'vtk' and pyvista.vtk_version_info < (9, 3):
+        raise ValueError("'vtk' pickle format requires VTK >= 9.3")
+
     pyvista.PICKLE_FORMAT = format_
 
 


Test patch:
diff --git a/tests/core/test_dataobject.py b/tests/core/test_dataobject.py
index 50b1308d..0b3485d1 100644
--- a/tests/core/test_dataobject.py
+++ b/tests/core/test_dataobject.py
@@ -214,8 +214,23 @@ def test_user_dict_persists_with_cells_to_points(uniform):
     assert uniform.user_dict['name'] == 'image'
 
 
-@pytest.mark.parametrize('pickle_format', ['xml', 'legacy'])
+def test_default_pickle_format():
+    assert pv.PICKLE_FORMAT == 'vtk' if pv.vtk_version_info >= (9, 3) else 'xml'
+
+
+@pytest.fixture
+def _modifies_pickle_format():
+    before = pv.PICKLE_FORMAT
+    yield
+    pv.PICKLE_FORMAT = before
+
+
+@pytest.mark.usefixtures('_modifies_pickle_format')
+@pytest.mark.parametrize('pickle_format', ['vtk', 'xml', 'legacy'])
 def test_pickle_serialize_deserialize(datasets, pickle_format):
+    if pickle_format == 'vtk' and pv.vtk_version_info < (9, 3):
+        pytest.xfail('VTK version not supported.')
+
     pv.set_pickle_format(pickle_format)
     for dataset in datasets:
         dataset_2 = pickle.loads(pickle.dumps(dataset))
@@ -256,11 +271,74 @@ def n_points(dataset):
     return dataset.n_points
 
 
-@pytest.mark.parametrize('pickle_format', ['xml', 'legacy'])
+@pytest.mark.usefixtures('_modifies_pickle_format')
+@pytest.mark.parametrize('pickle_format', ['vtk', 'xml', 'legacy'])
 def test_pickle_multiprocessing(datasets, pickle_format):
+    if pickle_format == 'vtk' and pv.vtk_version_info < (9, 3):
+        pytest.xfail('VTK version not supported.')
+
     # exercise pickling via multiprocessing
     pv.set_pickle_format(pickle_format)
     with multiprocessing.Pool(2) as p:
         res = p.map(n_points, datasets)
     for r, dataset in zip(res, datasets):
         assert r == dataset.n_points
+
+
+@pytest.mark.usefixtures('_modifies_pickle_format')
+@pytest.mark.parametrize('pickle_format', ['vtk', 'xml', 'legacy'])
+def test_pickle_multiblock(multiblock_all_with_nested_and_none, pickle_format):
+    if pickle_format == 'vtk' and pv.vtk_version_info < (9, 3):
+        pytest.xfail('VTK version not supported.')
+
+    pv.set_pickle_format(pickle_format)
+    multiblock = multiblock_all_with_nested_and_none
+
+    if pickle_format in ['legacy', 'xml']:
+        match = "MultiBlock is not supported with 'xml' or 'legacy' pickle formats.\nUse `pyvista.PICKLE_FORMAT='vtk'`."
+        with pytest.raises(TypeError, match=match):
+            pickle.dumps(multiblock)
+    else:
+        pickled = pickle.dumps(multiblock)
+        assert isinstance(pickled, bytes)
+        unpickled = pickle.loads(pickled)
+        assert unpickled == multiblock
+
+
+@pytest.mark.usefixtures('_modifies_pickle_format')
+@pytest.mark.parametrize('pickle_format', ['vtk', 'xml', 'legacy'])
+def test_pickle_user_dict(sphere, pickle_format):
+    if pickle_format == 'vtk' and pv.vtk_version_info < (9, 3):
+        pytest.xfail('VTK version not supported.')
+
+    pv.set_pickle_format(pickle_format)
+    user_dict = {'custom_attribute': 42}
+    sphere.user_dict = user_dict
+
+    pickled = pickle.dumps(sphere)
+    unpickled = pickle.loads(pickled)
+
+    assert unpickled.user_dict == user_dict
+
+
+@pytest.mark.usefixtures('_modifies_pickle_format')
+@pytest.mark.parametrize('pickle_format', ['vtk', 'xml', 'legacy'])
+def test_set_pickle_format(pickle_format):
+    if pickle_format == 'vtk' and pv.vtk_version_info < (9, 3):
+        match = 'requires VTK >= 9.3'
+        with pytest.raises(ValueError, match=match):
+            pv.set_pickle_format(pickle_format)
+    else:
+        pv.set_pickle_format(pickle_format)
+        assert pickle_format == pv.PICKLE_FORMAT
+
+
+@pytest.mark.usefixtures('_modifies_pickle_format')
+def test_pickle_invalid_format(sphere):
+    match = 'Unsupported pickle format `invalid_format`.'
+    with pytest.raises(ValueError, match=match):
+        pv.set_pickle_format('invalid_format')
+
+    pv.PICKLE_FORMAT = 'invalid_format'
+    with pytest.raises(ValueError, match=match):
+        pickle.dumps(sphere)
diff --git a/tests/core/test_utilities.py b/tests/core/test_utilities.py
index db2bbcdf..3cedb85c 100644
--- a/tests/core/test_utilities.py
+++ b/tests/core/test_utilities.py
@@ -794,17 +794,6 @@ def test_spherical_to_cartesian():
     assert np.allclose(pv.cartesian_to_spherical(x, y, z), points.T)
 
 
-def test_set_pickle_format():
-    pv.set_pickle_format('legacy')
-    assert pv.PICKLE_FORMAT == 'legacy'
-
-    pv.set_pickle_format('xml')
-    assert pv.PICKLE_FORMAT == 'xml'
-
-    with pytest.raises(ValueError):  # noqa: PT011
-        pv.set_pickle_format('invalid_format')
-
-
 def test_linkcode_resolve():
     assert linkcode_resolve('not-py', {}) is None
     link = linkcode_resolve('py', {'module': 'pyvista', 'fullname': 'pyvista.core.DataObject'})