Repository: googleapis/proto-plus-python

Issue:
Marshal fails with TypeError when using cross api message
See test code here https://github.com/googleapis/proto-plus-python/pull/348/commits/e123ad8813823c477d1ff1494a1c6c1ba85bff89, with corresponding failure [here](https://github.com/googleapis/proto-plus-python/actions/runs/4426034682/jobs/7761831001).

See the full stack trace below
```
________________________ test_module_package_cross_api _________________________

    def test_module_package_cross_api():
        sys.modules[__name__].__protobuf__ = proto.module(package="spam.eggs.v1")
        try:
    
            class Baz(proto.Message):
                foo = proto.RepeatedField(proto.INT64, number=1)
    
    
            marshal = proto.Marshal(name="spam.eggs.v1")
    
            assert Baz.meta.package == "spam.eggs.v1"
            assert Baz.pb() in marshal._rules
    
            sys.modules[__name__].__protobuf__ = proto.module(package="ham.pancakes.v1")
    
            class AnotherMessage(proto.Message):
                qux = proto.Field(proto.MESSAGE, number=1, message=Baz)
    
            marshal = proto.Marshal(name="ham.pancakes.v1")
    
            assert AnotherMessage.meta.package == "ham.pancakes.v1"
            assert AnotherMessage.pb() in marshal._rules
            # Confirm that Baz.pb() is no longer present in marshal._rules
            assert Baz.pb() not in marshal._rules
    
            # Test using multiple packages together
>           msg = AnotherMessage(qux=Baz())

tests/test_modules.py:65: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <[AttributeError('Unknown field for AnotherMessage: _pb') raised in repr()] AnotherMessage object at 0x7f112bbc07d0>
mapping = {'qux': }, ignore_unknown_fields = False, kwargs = {'qux': }
params = {'qux': }
marshal = <proto.marshal.marshal.Marshal object at 0x7f112bbc0910>, key = 'qux'
value = , pb_value = 

    def __init__(
        self,
        mapping=None,
        *,
        ignore_unknown_fields=False,
        **kwargs,
    ):
        # We accept several things for `mapping`:
        #   * An instance of this class.
        #   * An instance of the underlying protobuf descriptor class.
        #   * A dict
        #   * Nothing (keyword arguments only).
        if mapping is None:
            if not kwargs:
                # Special fast path for empty construction.
                super().__setattr__("_pb", self._meta.pb())
                return
    
            mapping = kwargs
        elif isinstance(mapping, self._meta.pb):
            # Make a copy of the mapping.
            # This is a constructor for a new object, so users will assume
            # that it will not have side effects on the arguments being
            # passed in.
            #
            # The `wrap` method on the metaclass is the public API for taking
            # ownership of the passed in protobuf object.
            mapping = copy.deepcopy(mapping)
            if kwargs:
                mapping.MergeFrom(self._meta.pb(**kwargs))
    
            super().__setattr__("_pb", mapping)
            return
        elif isinstance(mapping, type(self)):
            # Just use the above logic on mapping's underlying pb.
            self.__init__(mapping=mapping._pb, **kwargs)
            return
        elif isinstance(mapping, collections.abc.Mapping):
            # Can't have side effects on mapping.
            mapping = copy.copy(mapping)
            # kwargs entries take priority for duplicate keys.
            mapping.update(kwargs)
        else:
            # Sanity check: Did we get something not a map? Error if so.
            raise TypeError(
                "Invalid constructor input for %s: %r"
                % (
                    self.__class__.__name__,
                    mapping,
                )
            )
    
        params = {}
        # Update the mapping to address any values that need to be
        # coerced.
        marshal = self._meta.marshal
        for key, value in mapping.items():
            (key, pb_type) = self._get_pb_type_from_key(key)
            if pb_type is None:
                if ignore_unknown_fields:
                    continue
    
                raise ValueError(
                    "Unknown field for {}: {}".format(self.__class__.__name__, key)
                )
    
            try:
                pb_value = marshal.to_proto(pb_type, value)
            except ValueError:
                # Underscores may be appended to field names
                # that collide with python or proto-plus keywords.
                # In case a key only exists with a `_` suffix, coerce the key
                # to include the `_` suffix. It's not possible to
                # natively define the same field with a trailing underscore in protobuf.
                # See related issue
                # https://github.com/googleapis/python-api-core/issues/2[27](https://github.com/googleapis/proto-plus-python/actions/runs/4426034682/jobs/7761831001#step:6:28)
                if isinstance(value, dict):
                    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)
    
                pb_value = marshal.to_proto(pb_type, value)
    
            if pb_value is not None:
                params[key] = pb_value
    
        # Create the internal protocol buffer.
>       super().__setattr__("_pb", self._meta.pb(**params))
E       TypeError: Parameter to MergeFrom() must be instance of same class: expected spam.eggs.v1.Baz got Baz.

proto/message.py:[60](https://github.com/googleapis/proto-plus-python/actions/runs/4426034682/jobs/7761831001#step:6:61)4: TypeError
```

Discussion:
parthea: Waiting for approval on go/proto-plus-cross-api-issue

Patch:
diff --git a/proto/marshal/marshal.py b/proto/marshal/marshal.py
index bd21cbc..e9fbbb9 100644
--- a/proto/marshal/marshal.py
+++ b/proto/marshal/marshal.py
@@ -159,6 +159,22 @@ class BaseMarshal:
         for rule_class in stringy_numbers.STRINGY_NUMBER_RULES:
             self.register(rule_class._proto_type, rule_class())
 
+    def get_rule(self, proto_type):
+        # Rules are needed to convert values between proto-plus and pb.
+        # Retrieve the rule for the specified proto type.
+        # The NoopRule will be used when a rule is not found.
+        rule = self._rules.get(proto_type, self._noop)
+
+        # If we don't find a rule, also check under `_instances`
+        # in case there is a rule in another package.
+        # See https://github.com/googleapis/proto-plus-python/issues/349
+        if rule == self._noop and hasattr(self, "_instances"):
+            for _, instance in self._instances.items():
+                rule = instance._rules.get(proto_type, self._noop)
+                if rule != self._noop:
+                    break
+        return rule
+
     def to_python(self, proto_type, value, *, absent: bool = None):
         # Internal protobuf has its own special type for lists of values.
         # Return a view around it that implements MutableSequence.
@@ -174,10 +190,7 @@ class BaseMarshal:
         # Same thing for maps of messages.
         if value_type in compat.map_composite_types:
             return MapComposite(value, marshal=self)
-
-        # Convert ordinary values.
-        rule = self._rules.get(proto_type, self._noop)
-        return rule.to_python(value, absent=absent)
+        return self.get_rule(proto_type=proto_type).to_python(value, absent=absent)
 
     def to_proto(self, proto_type, value, *, strict: bool = False):
         # The protos in google/protobuf/struct.proto are exceptional cases,
@@ -212,9 +225,7 @@ class BaseMarshal:
             recursive_type = type(proto_type().value)
             return {k: self.to_proto(recursive_type, v) for k, v in value.items()}
 
-        # Convert ordinary values.
-        rule = self._rules.get(proto_type, self._noop)
-        pb_value = rule.to_proto(value)
+        pb_value = self.get_rule(proto_type=proto_type).to_proto(value)
 
         # Sanity check: If we are in strict mode, did we get the value we want?
         if strict and not isinstance(pb_value, proto_type):


Test patch:
diff --git a/tests/test_modules.py b/tests/test_modules.py
index 79a99da..7ab0e88 100644
--- a/tests/test_modules.py
+++ b/tests/test_modules.py
@@ -36,6 +36,38 @@ def test_module_package():
         del sys.modules[__name__].__protobuf__
 
 
+def test_module_package_cross_api():
+    sys.modules[__name__].__protobuf__ = proto.module(package="spam.eggs.v1")
+    try:
+
+        class Baz(proto.Message):
+            foo = proto.RepeatedField(proto.INT64, number=1)
+
+        marshal = proto.Marshal(name="spam.eggs.v1")
+
+        assert Baz.meta.package == "spam.eggs.v1"
+        assert Baz.pb() in marshal._rules
+
+        sys.modules[__name__].__protobuf__ = proto.module(package="ham.pancakes.v1")
+
+        class AnotherMessage(proto.Message):
+            qux = proto.Field(proto.MESSAGE, number=1, message=Baz)
+
+        marshal = proto.Marshal(name="ham.pancakes.v1")
+
+        assert AnotherMessage.meta.package == "ham.pancakes.v1"
+        assert AnotherMessage.pb() in marshal._rules
+        # Confirm that Baz.pb() is no longer present in marshal._rules
+        assert Baz.pb() not in marshal._rules
+
+        # Test using multiple packages together
+        # See https://github.com/googleapis/proto-plus-python/issues/349.
+        msg = AnotherMessage(qux=Baz())
+        assert type(msg) == AnotherMessage
+    finally:
+        del sys.modules[__name__].__protobuf__
+
+
 def test_module_package_explicit_marshal():
     sys.modules[__name__].__protobuf__ = proto.module(
         package="spam.eggs.v1",