fix: restore messages not attached to rpc in selective_gapic_generation

Author: parthea

packages/gapic-generator/gapic/schema/api.py

@@ -258,112 +258,21 @@ def disambiguate(self, string: str) -> str:
             return self.disambiguate(f"_{string}")
         return string
 
-    def add_to_address_allowlist(
+    def with_selective_generation(
         self,
         *,
-        address_allowlist: Set["metadata.Address"],
-        method_allowlist: Set[str],
-        resource_messages: Dict[str, "wrappers.MessageType"],
-    ) -> None:
-        """Adds to the set of Addresses of wrapper objects to be included in selective GAPIC generation.
-
-        This method is used to create an allowlist of addresses to be used to filter out unneeded
-        services, methods, messages, and enums at a later step.
-
-        Args:
-            address_allowlist (Set[metadata.Address]): A set of allowlisted metadata.Address
-                objects to add to. Only the addresses of the allowlisted methods, the services
-                containing these methods, and messages/enums those methods use will be part of the
-                final address_allowlist. The set may be modified during this call.
-            method_allowlist (Set[str]): An allowlist of fully-qualified method names.
-            resource_messages (Dict[str, wrappers.MessageType]): A dictionary mapping the unified
-                resource type name of a resource message to the corresponding MessageType object
-                representing that resource message. Only resources with a message representation
-                should be included in the dictionary.
-        Returns:
-            None
-        """
-        # The method.operation_service for an extended LRO is not fully qualified, so we
-        # truncate the service names accordingly so they can be found in
-        # method.add_to_address_allowlist
-        services_in_proto = {
-            service.name: service for service in self.services.values()
-        }
-        for service in self.services.values():
-            service.add_to_address_allowlist(
-                address_allowlist=address_allowlist,
-                method_allowlist=method_allowlist,
-                resource_messages=resource_messages,
-                services_in_proto=services_in_proto,
-            )
-
-    def prune_messages_for_selective_generation(
-        self, *, address_allowlist: Set["metadata.Address"]
-    ) -> Optional["Proto"]:
-        """Returns a truncated version of this Proto.
-
-        Only the services, messages, and enums contained in the allowlist
-        of visited addresses are included in the returned object. If there
-        are no services, messages, or enums left, and no file level resources,
-        return None.
-
-        Args:
-            address_allowlist (Set[metadata.Address]): A set of allowlisted metadata.Address
-                objects to filter against. Objects with addresses not the allowlist will be
-                removed from the returned Proto.
-        Returns:
-            Optional[Proto]: A truncated version of this proto. If there are no services, messages,
-                or enums left after the truncation process and there are no file level resources,
-                returns None.
-        """
-        # Once the address allowlist has been created, it suffices to only
-        # prune items at 2 different levels to truncate the Proto object:
-        #
-        #   1. At the Proto level, we remove unnecessary services, messages,
-        #      and enums.
-        #   2. For allowlisted services, at the Service level, we remove
-        #      non-allowlisted methods.
-        services = {
-            k: v.prune_messages_for_selective_generation(
-                address_allowlist=address_allowlist
-            )
-            for k, v in self.services.items()
-            if v.meta.address in address_allowlist
-        }
-
-        all_messages = {
-            k: v for k, v in self.all_messages.items() if v.ident in address_allowlist
-        }
-
-        all_enums = {
-            k: v for k, v in self.all_enums.items() if v.ident in address_allowlist
-        }
-
-        if not services and not all_messages and not all_enums:
-            return None
-
-        return dataclasses.replace(
-            self, services=services, all_messages=all_messages, all_enums=all_enums
-        )
-
-    def with_internal_methods(self, *, public_methods: Set[str]) -> "Proto":
-        """Returns a version of this Proto with some Methods marked as internal.
+        generate_omitted_as_internal: bool,
+        public_methods: Set[str],
+    ) -> "Proto":
 
-        The methods not in the public_methods set will be marked as internal and
-        services containing these methods will also be marked as internal by extension.
-        (See :meth:`Service.is_internal` for more details).
+        services = {}
+        for k, v in self.services.items():
+            new_v = v.with_selective_generation(
+                generate_omitted_as_internal=generate_omitted_as_internal,
+                public_methods=public_methods)
+            if new_v:
+                services[k] = new_v
 
-        Args:
-            public_methods (Set[str]): An allowlist of fully-qualified method names.
-                Methods not in this allowlist will be marked as internal.
-        Returns:
-            Proto: A version of this Proto with Method objects corresponding to methods
-                not in `public_methods` marked as internal.
-        """
-        services = {
-            k: v.with_internal_methods(public_methods=public_methods)
-            for k, v in self.services.items()
-        }
         return dataclasses.replace(self, services=services)
 
 
@@ -529,37 +438,13 @@ def disambiguate_keyword_sanitize_fname(
                     k: v for k, v in api.all_protos.items() if k not in api.protos
                 }
 
-                if selective_gapic_settings.generate_omitted_as_internal:
-                    for name, proto in api.protos.items():
-                        new_all_protos[name] = proto.with_internal_methods(
-                            public_methods=selective_gapic_methods
-                        )
-                else:
-                    all_resource_messages = collections.ChainMap(
-                        *(proto.resource_messages for proto in protos.values())
+                for name, proto in api.protos.items():
+                    proto_to_generate = proto.with_selective_generation(
+                        generate_omitted_as_internal=selective_gapic_settings.generate_omitted_as_internal,
+                        public_methods=selective_gapic_methods,
                     )
-
-                    # Prepare a list of addresses to include in selective generation,
-                    # then prune each Proto object. We look at metadata.Addresses, not objects, because
-                    # objects that refer to the same thing in the proto are different Python objects
-                    # in memory.
-                    address_allowlist: Set["metadata.Address"] = set([])
-                    for proto in api.protos.values():
-                        proto.add_to_address_allowlist(
-                            address_allowlist=address_allowlist,
-                            method_allowlist=selective_gapic_methods,
-                            resource_messages=all_resource_messages,
-                        )
-
-                    # We only prune services/messages/enums from protos that are not dependencies.
-                    for name, proto in api.protos.items():
-                        proto_to_generate = (
-                            proto.prune_messages_for_selective_generation(
-                                address_allowlist=address_allowlist
-                            )
-                        )
-                        if proto_to_generate:
-                            new_all_protos[name] = proto_to_generate
+                    if proto_to_generate:
+                        new_all_protos[name] = proto_to_generate
 
                 api = cls(
                     naming=naming,