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,7 +28,7 @@
 
 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
 
@@ -359,19 +359,27 @@ def resolve(self, selector: str) -> str:
             return f'{".".join(self.package)}.{selector}'
         return selector
 
-    def with_context(self, *, collisions: Set[str]) -> "Address":
+    def with_context(self, *, collisions: Set[str], context_cache: Optional[Dict[tuple, "MessageType"]] = 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 (
+        context_cache = context_cache if context_cache is not None else {}
+        cache_key = (id(self), id(collisions))
+        if cache_key in context_cache:
+            return context_cache[cache_key]
+    
+        updated_msg = (
             dataclasses.replace(self, collisions=collisions)
             if collisions and collisions != self.collisions
             else self
         )
 
+        context_cache[cache_key] = updated_msg
+        return updated_msg
+
 
 @dataclasses.dataclass(frozen=True)
 class Metadata:
@@ -398,14 +406,20 @@ def doc(self):
             return "\n\n".join(self.documentation.leading_detached_comments)
         return ""
 
-    def with_context(self, *, collisions: Set[str]) -> "Metadata":
+    def with_context(self, *, collisions: Set[str], context_cache: Optional[Dict[tuple, "MessageType"]] = 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 (
+        context_cache = context_cache if context_cache is not None else {}
+
+        cache_key = (id(self), id(collisions))
+        if cache_key in context_cache:
+            return context_cache[cache_key]
+        
+        updated_msg =  (
             dataclasses.replace(
                 self,
                 address=self.address.with_context(collisions=collisions),
@@ -414,6 +428,9 @@ def with_context(self, *, collisions: Set[str]) -> "Metadata":
             else self
         )
 
+        context_cache[cache_key] = updated_msg
+        return updated_msg
+
 
 @dataclasses.dataclass(frozen=True)
 class FieldIdentifier: