docarray · samsja · Feb 3, 2023 · Feb 3, 2023 · Feb 3, 2023 · Feb 3, 2023
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -133,6 +133,33 @@ jobs:
 #          fail_ci_if_error: false
 #          token: ${{ secrets.CODECOV_TOKEN }} # not required for public repos
 
+  docarray-test-proto3:
+    needs: [lint-ruff, check-black, import-test]
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.7]
+    steps:
+      - uses: actions/checkout@v2.5.0
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Prepare environment
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install poetry
+          poetry install --all-extras
+          pip install protobuf==3.19.0 # we check that we support 3.19 
+
+      - name: Test
+        id: test
+        run: |
+          poetry run pytest -k 'proto' tests
+        timeout-minutes: 30
+
+
   # just for blocking the merge until all parallel core-test are successful
   success-all-test:
     needs: [docarray-test, check-mypy]

diff --git a/docarray/proto/__init__.py b/docarray/proto/__init__.py
@@ -1,11 +1,23 @@
-from docarray.proto.pb2.docarray_pb2 import (
-    DocumentArrayProto,
-    DocumentArrayStackedProto,
-    DocumentProto,
-    NdArrayProto,
-    NodeProto,
-    UnionArrayProto,
-)
+from google.protobuf import __version__ as __pb__version__
+
+if __pb__version__.startswith('4'):
+    from docarray.proto.pb.docarray_pb2 import (
+        DocumentArrayProto,
+        DocumentArrayStackedProto,
+        DocumentProto,
+        NdArrayProto,
+        NodeProto,
+        UnionArrayProto,
+    )
+else:
+    from docarray.proto.pb2.docarray_pb2 import (
+        DocumentArrayProto,
+        DocumentArrayStackedProto,
+        DocumentProto,
+        NdArrayProto,
+        NodeProto,
+        UnionArrayProto,
+    )
 
 __all__ = [
     'DocumentArrayProto',

diff --git a/docarray/proto/pb/docarray_pb2.py b/docarray/proto/pb/docarray_pb2.py
diff --git a/docarray/proto/pb2/docarray_pb2.py b/docarray/proto/pb2/docarray_pb2.py
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,7 @@ license='Apache 2.0'
 python = "^3.7"
 pydantic = "^1.10.2"
 numpy = "^1.17.3"
-protobuf = { version = "^4.21.9", optional = true }
+protobuf = { version = ">=3.19.0", optional = true }
 torch = { version = "^1.0.0", optional = true }
 orjson = "^3.8.2"
 pillow = {version = "^9.3.0", optional = true }
@@ -71,15 +71,15 @@ exclude = 'docarray/proto/pb2/*'
 
 
 [tool.isort]
-skip_glob= ['docarray/proto/pb2/*']
+skip_glob= ['docarray/proto/pb2/*', 'docarray/proto/pb/*']
 
 [tool.ruff]
-exclude = ['docarray/proto/pb2/*', 'docs/*']
+exclude = ['docarray/proto/pb2/*', 'docarray/proto/pb/*', 'docs/*']
 
 [tool.pytest.ini_options]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
     "internet: marks tests as requiring internet (deselect with '-m \"not internet\"')",
     "asyncio: marks that run async tests",
-
+    "proto: mark tests that run with proto"
 ]
diff --git a/tests/integrations/document/test_proto.py b/tests/integrations/document/test_proto.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 import torch
 
 from docarray import BaseDocument, DocumentArray
@@ -19,6 +20,7 @@
 from docarray.typing.tensor import NdArrayEmbedding
 
 
+@pytest.mark.proto
 def test_multi_modal_doc_proto():
     class MyMultiModalDoc(BaseDocument):
         image: Image
@@ -31,6 +33,7 @@ class MyMultiModalDoc(BaseDocument):
     MyMultiModalDoc.from_protobuf(doc.to_protobuf())
 
 
+@pytest.mark.proto
 def test_all_types():
     class NestedDoc(BaseDocument):
         tensor: NdArray

diff --git a/tests/units/array/test_array_proto.py b/tests/units/array/test_array_proto.py
@@ -1,11 +1,13 @@
 import numpy as np
+import pytest
 
 from docarray import BaseDocument, DocumentArray
 from docarray.array.array_stacked import DocumentArrayStacked
 from docarray.documents import Image, Text
 from docarray.typing import NdArray
 
 
+@pytest.mark.proto
 def test_simple_proto():
     class CustomDoc(BaseDocument):
         text: str
@@ -22,6 +24,7 @@ class CustomDoc(BaseDocument):
         assert (doc1.tensor == doc2.tensor).all()
 
 
+@pytest.mark.proto
 def test_nested_proto():
     class CustomDocument(BaseDocument):
         text: Text
@@ -39,6 +42,7 @@ class CustomDocument(BaseDocument):
     DocumentArray[CustomDocument].from_protobuf(da.to_protobuf())
 
 
+@pytest.mark.proto
 def test_nested_proto_any_doc():
     class CustomDocument(BaseDocument):
         text: Text
@@ -56,6 +60,7 @@ class CustomDocument(BaseDocument):
     DocumentArray.from_protobuf(da.to_protobuf())
 
 
+@pytest.mark.proto
 def test_stacked_proto():
     class CustomDocument(BaseDocument):
         image: NdArray

diff --git a/tests/units/array/test_array_stacked.py b/tests/units/array/test_array_stacked.py
@@ -37,14 +37,12 @@ def test_iterator(batch):
 
 
 def test_stack_setter(batch):
-
     batch.tensor = torch.ones(10, 3, 224, 224)
 
     assert (batch.tensor == torch.ones(10, 3, 224, 224)).all()
 
 
 def test_stack_optional(batch):
-
     assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
     assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()
 
@@ -67,7 +65,6 @@ class Image(BaseDocument):
 
 
 def test_stack(batch):
-
     assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
     assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()
     assert batch._columns['tensor'].data_ptr() == batch.tensor.data_ptr()
@@ -138,11 +135,12 @@ class MMdoc(BaseDocument):
         assert (doc.img.tensor == torch.zeros(3, 224, 224)).all()
 
 
+@pytest.mark.proto
 def test_proto_stacked_mode_torch(batch):
-
     batch.from_protobuf(batch.to_protobuf())
 
 
+@pytest.mark.proto
 def test_proto_stacked_mode_numpy():
     class MyDoc(BaseDocument):
         tensor: NdArray[3, 224, 224]