fix: restore messages not attached to rpc in selective_gapic_generation

[parthea]

Towards b/507893758, b/507889482, b/501132869, b/503326310

[gemini-code-assist]

Code Review

This pull request refactors the selective GAPIC generation logic by replacing the multi-step allowlist and pruning process with a unified with_selective_generation method across the Proto, Service, and Method classes. This consolidation simplifies the code by handling both internal method marking and service/method omission in a single pass. The review feedback highlights the need for docstrings on the new methods and recommends updating return type hints to Optional to correctly handle cases where objects are filtered out (returning None). Specifically, the Proto implementation should be updated to return None if it becomes empty after filtering.

[gemini-code-assist]

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.

[gemini-code-assist]

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.

[gemini-code-assist]

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.