optimize gapic

Author: ohmayr

gapic/schema/api.py

@@ -115,6 +115,7 @@ def build(
         prior_protos: Optional[Mapping[str, "Proto"]] = None,
         load_services: bool = True,
         all_resources: Optional[Mapping[str, wrappers.MessageType]] = None,
+        context_cache: Optional[Dict[tuple, "wrappers.MessageType"]] = None,
     ) -> "Proto":
         """Build and return a Proto instance.
 
@@ -139,6 +140,7 @@ def build(
             prior_protos=prior_protos or {},
             load_services=load_services,
             all_resources=all_resources or {},
+            context_cache=context_cache,
         ).proto
 
     @cached_property
@@ -455,6 +457,7 @@ def disambiguate_keyword_sanitize_fname(
         # We just load all the APIs types first and then
         # load the services and methods with the full scope of types.
         pre_protos: Dict[str, Proto] = dict(prior_protos or {})
+        context_cache = {}
         for fd in file_descriptors:
             fd.name = disambiguate_keyword_sanitize_fname(fd.name, pre_protos)
             pre_protos[fd.name] = Proto.build(
@@ -465,6 +468,7 @@ def disambiguate_keyword_sanitize_fname(
                 prior_protos=pre_protos,
                 # Ugly, ugly hack.
                 load_services=False,
+                context_cache=context_cache,
             )
 
         # A file descriptor's file-level resources are NOT visible to any importers.
@@ -485,6 +489,7 @@ def disambiguate_keyword_sanitize_fname(
                 opts=opts,
                 prior_protos=pre_protos,
                 all_resources=MappingProxyType(all_file_resources),
+                context_cache=context_cache,
             )
             for name, proto in pre_protos.items()
         }
@@ -1103,6 +1108,8 @@ def __init__(
         prior_protos: Optional[Mapping[str, Proto]] = None,
         load_services: bool = True,
         all_resources: Optional[Mapping[str, wrappers.MessageType]] = None,
+        context_cache: Optional[Dict[tuple, "wrappers.MessageType"]] = None,
+        
     ):
         self.proto_messages: Dict[str, wrappers.MessageType] = {}
         self.proto_enums: Dict[str, wrappers.EnumType] = {}
@@ -1111,6 +1118,7 @@ def __init__(
         self.file_to_generate = file_to_generate
         self.prior_protos = prior_protos or {}
         self.opts = opts
+        self.context_cache = context_cache
 
         # Iterate over the documentation and place it into a dictionary.
         #
@@ -1220,20 +1228,22 @@ def proto(self) -> Proto:
         if not self.file_to_generate:
             return naive
 
+        global_collisions = frozenset(naive.names)
         visited_messages: Set[wrappers.MessageType] = set()
         # Return a context-aware proto object.
         return dataclasses.replace(
             naive,
             all_enums=collections.OrderedDict(
-                (k, v.with_context(collisions=naive.names))
+                (k, v.with_context(collisions=global_collisions, context_cache=self.context_cache))
                 for k, v in naive.all_enums.items()
             ),
             all_messages=collections.OrderedDict(
                 (
                     k,
                     v.with_context(
-                        collisions=naive.names,
+                        collisions=global_collisions,
                         visited_messages=visited_messages,
+                        context_cache=self.context_cache,
                     ),
                 )
                 for k, v in naive.all_messages.items()
@@ -1244,8 +1254,9 @@ def proto(self) -> Proto:
                 (
                     k,
                     v.with_context(
-                        collisions=v.names,
-                        visited_messages=visited_messages,
+                        collisions=global_collisions,
+                        visited_messages=frozenset(v.names),
+                        context_cache=self.context_cache,
                     ),
                 )
                 for k, v in naive.services.items()

gapic/schema/metadata.py

@@ -28,13 +28,14 @@
 
 import dataclasses
 import re
-from typing import FrozenSet, Set, Tuple, Optional
+from typing import FrozenSet, Set, Tuple, Optional, Dict
 
 from google.protobuf import descriptor_pb2
 
 from gapic.schema import imp
 from gapic.schema import naming
 from gapic.utils import cached_property
+from gapic.utils import cached_proto_context
 from gapic.utils import RESERVED_NAMES
 
 # This class is a minor hack to optimize Address's __eq__ method.
@@ -359,19 +360,23 @@ def resolve(self, selector: str) -> str:
             return f'{".".join(self.package)}.{selector}'
         return selector
 
-    def with_context(self, *, collisions: Set[str]) -> "Address":
+    @cached_proto_context
+    def with_context(self, *, collisions: Set[str], context_cache: Optional[Dict] = None) -> "Address":
         """Return a derivative of this address with the provided context.
 
         This method is used to address naming collisions. The returned
         ``Address`` object aliases module names to avoid naming collisions in
         the file being written.
         """
-        return (
+    
+        updated_msg = (
             dataclasses.replace(self, collisions=collisions)
             if collisions and collisions != self.collisions
             else self
         )
 
+        return updated_msg
+
 
 @dataclasses.dataclass(frozen=True)
 class Metadata:
@@ -398,22 +403,26 @@ def doc(self):
             return "\n\n".join(self.documentation.leading_detached_comments)
         return ""
 
-    def with_context(self, *, collisions: Set[str]) -> "Metadata":
+    @cached_proto_context
+    def with_context(self, *, collisions: Set[str], context_cache: Optional[Dict] = None) -> "Metadata":
         """Return a derivative of this metadata with the provided context.
 
         This method is used to address naming collisions. The returned
         ``Address`` object aliases module names to avoid naming collisions in
         the file being written.
         """
-        return (
+        
+        updated_msg =  (
             dataclasses.replace(
                 self,
-                address=self.address.with_context(collisions=collisions),
+                address=self.address.with_context(collisions=collisions, context_cache=context_cache),
             )
             if collisions and collisions != self.address.collisions
             else self
         )
 
+        return updated_msg
+
 
 @dataclasses.dataclass(frozen=True)
 class FieldIdentifier: