feat: add initial files to allow protobuf 4.20 (#327)

* feat: add initial files to allow protobuf 4.20
This PR integrates protobuf 4.xx into codebase.

Co-authored-by: Anthonios Partheniou <[email protected]>

* chore: address pr comment.

Co-authored-by: Anthonios Partheniou <[email protected]>

.github/workflows/tests.yml

@@ -47,7 +47,7 @@ jobs:
     strategy:
       matrix:
         python: ['3.6', '3.7', '3.8', '3.9', '3.10']
-        variant: ['', cpp]
+        variant: ['', 'cpp', 'upb']
     steps:
       - name: Cancel Previous Runs
         uses: styfle/[email protected]

noxfile.py

@@ -42,7 +42,8 @@ def unit(session, proto="python"):
     session.env["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = proto
     session.install("coverage", "pytest", "pytest-cov", "pytz")
     session.install("-e", ".[testing]", "-c", constraints_path)
-
+    if proto == "cpp":  # 4.20 does not have cpp.
+        session.install("protobuf==3.19.0")
     session.run(
         "py.test",
         "-W=error",
@@ -54,7 +55,7 @@ def unit(session, proto="python"):
                 "--cov-config=.coveragerc",
                 "--cov-report=term",
                 "--cov-report=html",
-                os.path.join("tests", ""),
+                "tests",
             ]
         ),
     )
@@ -68,6 +69,11 @@ def unitcpp(session):
     return unit(session, proto="cpp")
 
 
+@nox.session(python=PYTHON_VERSIONS)
+def unitupb(session):
+    return unit(session, proto="upb")
+
+
 # Just use the most recent version for docs
 @nox.session(python=PYTHON_VERSIONS[-1])
 def docs(session):

proto/_package_info.py

@@ -37,7 +37,13 @@ def compile(name, attrs):
     proto_module = getattr(module, "__protobuf__", object())
 
     # A package should be present; get the marshal from there.
-    package = getattr(proto_module, "package", module_name)
+    # TODO: Revert to empty string as a package value after protobuf fix.
+    # When package is empty, upb based protobuf fails with an
+    # "TypeError: Couldn't build proto file into descriptor pool: invalid name: empty part ()' means"
+    # during an attempt to add to descriptor pool.
+    package = getattr(
+        proto_module, "package", module_name if module_name else "_default_package"
+    )
     marshal = Marshal(name=getattr(proto_module, "marshal", package))
 
     # Done; return the data.

proto/marshal/collections/maps.py

@@ -15,6 +15,7 @@
 import collections
 
 from proto.utils import cached_property
+from google.protobuf.message import Message
 
 
 class MapComposite(collections.abc.MutableMapping):
@@ -58,15 +59,21 @@ def __getitem__(self, key):
 
     def __setitem__(self, key, value):
         pb_value = self._marshal.to_proto(self._pb_type, value, strict=True)
-
         # Directly setting a key is not allowed; however, protocol buffers
         # is so permissive that querying for the existence of a key will in
         # of itself create it.
         #
         # Therefore, we create a key that way (clearing any fields that may
         # be set) and then merge in our values.
-        self.pb[key].Clear()
-        self.pb[key].MergeFrom(pb_value)
+        # TODO: self.pb[key] should always be Message. Remove this after protobuf fix.
+        # In UPB, sometimes self.pb[key] is not always a proto.
+        # This happens during marshalling when the pb_value is upb.MapCompositeContainer
+        # so it's not marshalled correcrtly (i.e. should be scalar values not composite).
+        if isinstance(self.pb[key], Message):
+            self.pb[key].Clear()
+            self.pb[key].MergeFrom(pb_value)
+        else:
+            self.pb[key] = value
 
     def __delitem__(self, key):
         self.pb.pop(key)

proto/marshal/compat.py

@@ -20,14 +20,23 @@
 
 from google.protobuf.internal import containers
 
+# Import protobuf 4.xx first and fallback to earlier version
+# if not present.
 try:
-    from google.protobuf.pyext import _message
+    from google._upb import _message
 except ImportError:
     _message = None
 
+if not _message:
+    try:
+        from google.protobuf.pyext import _message
+    except ImportError:
+        _message = None
+
 repeated_composite_types = (containers.RepeatedCompositeFieldContainer,)
 repeated_scalar_types = (containers.RepeatedScalarFieldContainer,)
 map_composite_types = (containers.MessageMap,)
+
 if _message:
     repeated_composite_types += (_message.RepeatedCompositeContainer,)
     repeated_scalar_types += (_message.RepeatedScalarContainer,)

proto/marshal/marshal.py

@@ -25,6 +25,7 @@
 from proto.marshal.collections import MapComposite
 from proto.marshal.collections import Repeated
 from proto.marshal.collections import RepeatedComposite
+
 from proto.marshal.rules import bytes as pb_bytes
 from proto.marshal.rules import stringy_numbers
 from proto.marshal.rules import dates
@@ -219,7 +220,6 @@ def to_proto(self, proto_type, value, *, strict: bool = False):
                     got=pb_value.__class__.__name__,
                 ),
             )
-
         # Return the final value.
         return pb_value
 

proto/message.py

@@ -29,6 +29,10 @@
 from proto.fields import RepeatedField
 from proto.marshal import Marshal
 from proto.primitives import ProtoType
+from proto.utils import has_upb
+
+
+_upb = has_upb()  # Important to cache result here.
 
 
 class MessageMeta(type):
@@ -568,11 +572,21 @@ def __init__(
                 # See related issue
                 # https://github.com/googleapis/python-api-core/issues/227
                 if isinstance(value, dict):
-                    keys_to_update = [
-                        item
-                        for item in value
-                        if not hasattr(pb_type, item) and hasattr(pb_type, f"{item}_")
-                    ]
+                    if _upb:
+                        # In UPB, pb_type is MessageMeta which doesn't expose attrs like it used to in Python/CPP.
+                        keys_to_update = [
+                            item
+                            for item in value
+                            if item not in pb_type.DESCRIPTOR.fields_by_name
+                            and f"{item}_" in pb_type.DESCRIPTOR.fields_by_name
+                        ]
+                    else:
+                        keys_to_update = [
+                            item
+                            for item in value
+                            if not hasattr(pb_type, item)
+                            and hasattr(pb_type, f"{item}_")
+                        ]
                     for item in keys_to_update:
                         value[f"{item}_"] = value.pop(item)
 

proto/utils.py

@@ -15,6 +15,16 @@
 import functools
 
 
+def has_upb():
+    try:
+        from google._upb import _message  # pylint: disable=unused-import
+
+        has_upb = True
+    except ImportError:
+        has_upb = False
+    return has_upb
+
+
 def cached_property(fx):
     """Make the callable into a cached property.
 

tests/conftest.py

@@ -23,6 +23,7 @@
 from proto._file_info import _FileInfo
 from proto.marshal import Marshal
 from proto.marshal import rules
+from proto.utils import has_upb
 
 
 def pytest_runtest_setup(item):
@@ -37,10 +38,14 @@ def pytest_runtest_setup(item):
         mock.patch.object(symbol_database, "Default", return_value=sym_db),
     ]
     if descriptor_pool._USE_C_DESCRIPTORS:
-        from google.protobuf.pyext import _message
 
         item._mocks.append(
-            mock.patch("google.protobuf.pyext._message.default_pool", pool)
+            mock.patch(
+                "google._upb._message.default_pool"
+                if has_upb()
+                else "google.protobuf.pyext._message.default_pool",
+                pool,
+            )
         )
 
     [i.start() for i in item._mocks]

tests/test_fields_repeated_composite.py

@@ -35,6 +35,7 @@ class Baz(proto.Message):
     assert len(baz.foos) == 1
     assert baz.foos == baz.foos
     assert baz.foos[0].bar == 42
+    assert isinstance(baz.foos[0], Foo)
 
 
 def test_repeated_composite_equality():

tests/test_message_filename.py

@@ -20,4 +20,7 @@ def test_filename_includes_classname_salt():
     class Foo(proto.Message):
         bar = proto.Field(proto.INT32, number=1)
 
-    assert Foo.pb(Foo()).DESCRIPTOR.file.name == "test_message_filename_foo.proto"
+    assert (
+        Foo.pb(Foo()).DESCRIPTOR.file.name
+        == "test_message_filename__default_package.foo.proto"
+    )