Address review feedback

Author: parthea

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

@@ -264,8 +264,18 @@ def with_selective_generation(
         generate_omitted_as_internal: bool,
         public_methods: Set[str],
         excluded_addresses: Set["metadata.Address"],
-    ) -> "Proto":
+    ) -> 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.
+            excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
 
+        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(
@@ -288,7 +298,7 @@ def with_selective_generation(
 
         # 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:
+        if not generate_omitted_as_internal and not services and not all_messages and not all_enums:
             return None
 
         return dataclasses.replace(

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

@@ -2066,7 +2066,7 @@ def with_selective_generation(
         generate_omitted_as_internal: bool,
         public_methods: Set[str],
         excluded_addresses: Set["metadata.Address"],
-    ) -> "Method":
+    ) -> Optional["Method"]:
 
         if self.ident.proto in public_methods:
             return self
@@ -2474,7 +2474,7 @@ def with_selective_generation(
         generate_omitted_as_internal: bool,
         public_methods: Set[str],
         excluded_addresses: Set["metadata.Address"],
-    ) -> "Service":
+    ) -> Optional["Service"]:
 
         methods = {}
         for k, v in self.methods.items():