Skip to content

fix: restore messages not attached to rpc in selective_gapic_generation #16951

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
parthea wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

fix: restore messages not attached to rpc in selective_gapic_generation #16951

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5cdf67e
fix: restore messages not attached to rpc in selective_gapic_generation
parthea May 5, 2026
e1c4e11
refine prototype
parthea May 5, 2026
5d8ec63
update goldens
parthea May 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
192 changes: 74 additions & 118 deletions packages/gapic-generator/gapic/schema/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -258,114 +258,46 @@ 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
}
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,
excluded_addresses=excluded_addresses)
if new_v:
services[k] = new_v

# 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 in address_allowlist
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 in address_allowlist
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
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.

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).

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)


@dataclasses.dataclass(frozen=True)
class API:
Expand Down Expand Up @@ -529,37 +461,61 @@ 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
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,
)
else:
all_resource_messages = collections.ChainMap(
*(proto.resource_messages for proto in protos.values())
)

# 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,
# 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,
)

# 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
# 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

api = cls(
naming=naming,
Expand Down
90 changes: 35 additions & 55 deletions packages/gapic-generator/gapic/schema/wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2014,7 +2014,6 @@ def add_to_address_allowlist(
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
Expand Down Expand Up @@ -2061,26 +2060,28 @@ def add_to_address_allowlist(
resource_messages=resource_messages,
)

def with_internal_methods(self, *, public_methods: Set[str]) -> "Method":
"""Returns a version of this ``Method`` marked as internal

The methods not in the public_methods set will be marked as internal and
this ``Service`` will as well by extension (see :meth:`Service.is_internal`).
def with_selective_generation(
self,
*,
generate_omitted_as_internal: bool,
public_methods: Set[str],
excluded_addresses: Set["metadata.Address"],
) -> "Method":

Args:
public_methods (Set[str]): An allowlist of fully-qualified method names.
Methods not in this allowlist will be marked as internal.
Returns:
Service: A version of this `Service` with `Method` objects corresponding to methods
not in `public_methods` marked as internal.
"""
if self.ident.proto in public_methods:
return self

return dataclasses.replace(
self,
is_internal=True,
)
# Not public.
# We mark it as internal if either generate_omitted_as_internal is set,
# or if the method is reachable from some public method (e.g. as a polling method).
if generate_omitted_as_internal or self.meta.address not in excluded_addresses:
return dataclasses.replace(
self,
is_internal=True,
)
else:
return None
Comment on lines +2063 to +2083
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The with_selective_generation method is missing a docstring and the return type hint should be Optional["Method"] since it returns None when a method is omitted and not marked as internal.

    def with_selective_generation(
        self,
        *,
        generate_omitted_as_internal: bool,
        public_methods: Set[str],
    ) -> Optional["Method"]:
        """Returns a version of this Method for selective generation.

        Args:
            generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
            public_methods (Set[str]): The set of fully-qualified method names to keep as public.

        Returns:
            Optional[Method]: The method (possibly marked internal), or None if it should be removed.
        """
        if self.ident.proto in public_methods:
            return self

        # Not public
        if generate_omitted_as_internal:
            return dataclasses.replace(
                self,
                is_internal=True,
            )
        else:
            return None
References
  1. When finding a precise type hint that satisfies both mypy and unit tests is not cost-effective, using a less specific type (e.g., Any) is an acceptable trade-off, especially if it improves upon the previous state.




@dataclasses.dataclass(frozen=True)
Expand Down Expand Up @@ -2467,45 +2468,24 @@ def add_to_address_allowlist(
services_in_proto=services_in_proto,
)

def prune_messages_for_selective_generation(
self, *, address_allowlist: Set["metadata.Address"]
def with_selective_generation(
self,
*,
generate_omitted_as_internal: bool,
public_methods: Set[str],
excluded_addresses: Set["metadata.Address"],
) -> "Service":
"""Returns a truncated version of this Service.

Only the methods, messages, and enums contained in the address allowlist
are included in the returned object.

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:
Service: A truncated version of this proto.
"""
return dataclasses.replace(
self,
methods={
k: v for k, v in self.methods.items() if v.ident in address_allowlist
},
)

def with_internal_methods(self, *, public_methods: Set[str]) -> "Service":
"""Returns a version of this ``Service`` with some Methods marked as internal.
methods = {}
for k, v in self.methods.items():
new_v = v.with_selective_generation(
generate_omitted_as_internal=generate_omitted_as_internal,
public_methods=public_methods,
excluded_addresses=excluded_addresses)
if new_v:
methods[k] = new_v

The methods not in the public_methods set will be marked as internal and
this ``Service`` will as well by extension (see :meth:`Service.is_internal`).
if not generate_omitted_as_internal and not methods:
return None

Args:
public_methods (Set[str]): An allowlist of fully-qualified method names.
Methods not in this allowlist will be marked as internal.
Returns:
Service: A version of this `Service` with `Method` objects corresponding to methods
not in `public_methods` marked as internal.
"""
return dataclasses.replace(
self,
methods={
k: v.with_internal_methods(public_methods=public_methods)
for k, v in self.methods.items()
},
)
return dataclasses.replace(self, methods=methods)
Comment on lines +2471 to +2491
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The with_selective_generation method is missing a docstring and the return type hint should be Optional["Service"] since it returns None when a service has no methods left and generate_omitted_as_internal is False.

    def with_selective_generation(
        self,
        *,
        generate_omitted_as_internal: bool,
        public_methods: Set[str],
    ) -> Optional["Service"]:
        """Returns a version of this Service for selective generation.

        Args:
            generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
            public_methods (Set[str]): The set of fully-qualified method names to keep as public.

        Returns:
            Optional[Service]: The service with filtered methods, or None if it should be removed.
        """
        methods = {}
        for k, v in self.methods.items():
            new_v = v.with_selective_generation(
                generate_omitted_as_internal=generate_omitted_as_internal,
                public_methods=public_methods)
            if new_v:
                methods[k] = new_v

        if not generate_omitted_as_internal and not methods:
            return None

        return dataclasses.replace(self, methods=methods)
References
  1. When finding a precise type hint that satisfies both mypy and unit tests is not cost-effective, using a less specific type (e.g., Any) is an acceptable trade-off, especially if it improves upon the previous state.

4 changes: 4 additions & 0 deletions .../gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from google.cloud.redis_v1.types.cloud_redis import Instance
from google.cloud.redis_v1.types.cloud_redis import ListInstancesRequest
from google.cloud.redis_v1.types.cloud_redis import ListInstancesResponse
from google.cloud.redis_v1.types.cloud_redis import LocationMetadata
from google.cloud.redis_v1.types.cloud_redis import MaintenancePolicy
from google.cloud.redis_v1.types.cloud_redis import MaintenanceSchedule
from google.cloud.redis_v1.types.cloud_redis import NodeInfo
Expand All @@ -35,6 +36,7 @@
from google.cloud.redis_v1.types.cloud_redis import TlsCertificate
from google.cloud.redis_v1.types.cloud_redis import UpdateInstanceRequest
from google.cloud.redis_v1.types.cloud_redis import WeeklyMaintenanceWindow
from google.cloud.redis_v1.types.cloud_redis import ZoneMetadata

__all__ = ('CloudRedisClient',
'CloudRedisAsyncClient',
Expand All @@ -44,6 +46,7 @@
'Instance',
'ListInstancesRequest',
'ListInstancesResponse',
'LocationMetadata',
'MaintenancePolicy',
'MaintenanceSchedule',
'NodeInfo',
Expand All @@ -52,4 +55,5 @@
'TlsCertificate',
'UpdateInstanceRequest',
'WeeklyMaintenanceWindow',
'ZoneMetadata',
)
Loading
Loading