refine prototype

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,58 @@ 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())
+                )
+
+                # Calculate all reachable addresses (API-wide).
+                # This includes all messages and enums reachable from ANY RPC
+                # defined in any proto of the API.
+                all_rpc_addresses: Set["metadata.Address"] = set([])
+                all_methods = set(api.all_methods.keys())
+                # 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
+
+                for proto in api.all_protos.values():
+                    for service in proto.services.values():
+                        service.add_to_address_allowlist(
+                            address_allowlist=all_rpc_addresses,
+                            method_allowlist=all_methods,
+                            resource_messages=all_resource_messages,
+                            services_in_proto=all_services,
+                        )
+
+                # 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 reachable from SOME RPC but NOT from any PUBLIC RPC.
+                # Types not attached to any RPC will not be in all_rpc_addresses and thus
+                # will NOT be in excluded_addresses, meaning they are preserved.
+                excluded_addresses = (
+                    all_rpc_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