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 1 commit
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
149 changes: 17 additions & 132 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,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)
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 its return type hint should be Optional["Proto"] to allow skipping empty protos. This is consistent with the logic in Service.with_selective_generation and the check at line 446. Additionally, it should return None if no services or messages remain when not generating omitted methods as internal.

    def with_selective_generation(
        self,
        *,
        generate_omitted_as_internal: bool,
        public_methods: Set[str],
    ) -> Optional["Proto"]:
        """Returns a version of this Proto 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[Proto]: A version of this Proto with services/methods filtered.
                Returns None if the Proto becomes empty and generate_omitted_as_internal is False.
        """
        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

        if not generate_omitted_as_internal and not services and not self.all_messages and not self.all_enums:
            return None

        return dataclasses.replace(self, services=services)
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.



Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading