fix: allow nested model dump via docvec (#1808)

Signed-off-by: samsja <sami.jaghouar@hotmail.fr>
docarray · Sep 28, 2023 · 4a1bc26 · 4a1bc26
1 parent 26d776d
commit 4a1bc26
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/docarray/base_doc/doc.py b/docarray/base_doc/doc.py
@@ -419,7 +419,6 @@ def dict(
             which fields to include or exclude.
 
             """
-
             exclude, original_exclude, doclist_exclude_fields = self._exclude_doclist(
                 exclude=exclude
             )
@@ -447,6 +446,20 @@ def dict(
 
     else:
 
+        def _copy_view_pydantic_v2(self: T) -> T:
+            """
+            perform a deep copy, the new doc has its own data
+            """
+            data = {}
+            for key, value in self.__dict__.to_dict().items():
+                if isinstance(value, BaseDocWithoutId):
+                    data[key] = value._copy_view_pydantic_v2()
+                else:
+                    data[key] = value
+
+            doc = self.__class__.model_construct(**data)
+            return doc
+
         def model_dump(  # type: ignore
             self,
             *,
@@ -460,15 +473,15 @@ def model_dump(  # type: ignore
             round_trip: bool = False,
             warnings: bool = True,
         ) -> Dict[str, Any]:
-            def _model_dump(cls):
+            def _model_dump(doc):
 
                 (
                     exclude_,
                     original_exclude,
                     doclist_exclude_fields,
                 ) = self._exclude_doclist(exclude=exclude)
 
-                data = cls.model_dump(
+                data = doc.model_dump(
                     mode=mode,
                     include=include,
                     exclude=exclude_,
@@ -495,7 +508,7 @@ def _model_dump(cls):
                 ## for some reason use ColumnViewStorage to dump the data is not working with
                 ## pydantic v2, so we need to create a new doc and dump it
 
-                new_doc = self.__class__.model_construct(**self.__dict__.to_dict())
+                new_doc = self._copy_view_pydantic_v2()
                 return _model_dump(new_doc)
             else:
                 return _model_dump(super())

diff --git a/tests/units/array/test_array_from_to_pandas.py b/tests/units/array/test_array_from_to_pandas.py
@@ -22,7 +22,6 @@ class MyDocNested(MyDoc):
     return MyDocNested
 
 
-@pytest.mark.skipif(is_pydantic_v2, reason="Not working with pydantic v2")
 @pytest.mark.parametrize('doc_vec', [False, True])
 def test_to_from_pandas_df(nested_doc_cls, doc_vec):
     da = DocList[nested_doc_cls](