Prune top level request/response messages

Author: parthea

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

@@ -263,17 +263,40 @@ def with_selective_generation(
         *,
         generate_omitted_as_internal: bool,
         public_methods: Set[str],
+        excluded_addresses: Set["metadata.Address"],
     ) -> "Proto":
 
         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)
+                public_methods=public_methods,
+                excluded_addresses=excluded_addresses)
             if new_v:
                 services[k] = new_v
 
-        return dataclasses.replace(self, services=services)
+        # We only prune messages/enums from protos that are not dependencies.
+        # Messages and enums are excluded only if they are reachable from some RPC
+        # but NOT from any of the publicly allowed RPCs.
+        all_messages = {
+            k: v for k, v in self.all_messages.items() if v.ident not in excluded_addresses
+        }
+
+        all_enums = {
+            k: v for k, v in self.all_enums.items() if v.ident not in excluded_addresses
+        }
+
+        # If the proto becomes empty after pruning, we return None to signal
+        # that it should be excluded from generation.
+        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,
+        )
 
 
 @dataclasses.dataclass(frozen=True)
@@ -438,10 +461,64 @@ def disambiguate_keyword_sanitize_fname(
                     k: v for k, v in api.all_protos.items() if k not in api.protos
                 }
 
+                all_resource_messages = collections.ChainMap(
+                    *(proto.resource_messages for proto in api.all_protos.values())
+                )
+
+                # Create a global map of services to support cross-proto lookup
+                # for extended LROs.
+                all_services: Dict[str, wrappers.Service] = {}
+                for p in api.all_protos.values():
+                    for s in p.services.values():
+                        all_services[s.meta.address.proto] = s
+                        all_services[s.name] = s
+
+                # Calculate addresses of omitted RPCs and their top-level request/response messages.
+                # These are the ONLY things that can be excluded.
+                # Other messages reachable from these (but not from public RPCs) will be KEPT.
+                candidate_excluded_addresses: Set["metadata.Address"] = set([])
+                for proto in api.all_protos.values():
+                    for service in proto.services.values():
+                        for method in service.methods.values():
+                            if method.ident.proto not in selective_gapic_methods:
+                                candidate_excluded_addresses.add(method.meta.address)
+                                candidate_excluded_addresses.add(method.input.ident)
+                                candidate_excluded_addresses.add(method.output.ident)
+
+                                # If this is an LRO, add its response and metadata types to candidates.
+                                if method.lro:
+                                    candidate_excluded_addresses.add(method.lro.response_type.ident)
+                                    candidate_excluded_addresses.add(method.lro.metadata_type.ident)
+
+                                # If this is an extended LRO, add its request and operation types to candidates.
+                                if method.extended_lro:
+                                    candidate_excluded_addresses.add(method.extended_lro.request_type.ident)
+                                    candidate_excluded_addresses.add(method.extended_lro.operation_type.ident)
+
+                # Calculate publicly reachable addresses (API-wide).
+                # This includes only types reachable from the allowlisted methods.
+                public_rpc_addresses: Set["metadata.Address"] = set([])
+                for proto in api.all_protos.values():
+                    for service in proto.services.values():
+                        service.add_to_address_allowlist(
+                            address_allowlist=public_rpc_addresses,
+                            method_allowlist=selective_gapic_methods,
+                            resource_messages=all_resource_messages,
+                            services_in_proto=all_services,
+                        )
+
+                # Addresses to exclude: those that are candidates but NOT reachable from any PUBLIC RPC.
+                excluded_addresses = (
+                    candidate_excluded_addresses - public_rpc_addresses
+                    if not selective_gapic_settings.generate_omitted_as_internal
+                    else set([])
+                )
+
                 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,
+                        excluded_addresses=excluded_addresses,
                     )
                     if proto_to_generate:
                         new_all_protos[name] = proto_to_generate