use a context manager instead of passing args

Author: ohmayr

gapic/cli/generate.py

@@ -23,7 +23,7 @@
 from gapic import generator
 from gapic.schema import api
 from gapic.utils import Options
-
+from gapic.utils.cache import generation_cache_context
 
 @click.command()
 @click.option(
@@ -56,15 +56,16 @@ def generate(request: typing.BinaryIO, output: typing.BinaryIO) -> None:
         [p.package for p in req.proto_file if p.name in req.file_to_generate]
     ).rstrip(".")
 
-    # Build the API model object.
-    # This object is a frozen representation of the whole API, and is sent
-    # to each template in the rendering step.
-    api_schema = api.API.build(req.proto_file, opts=opts, package=package)
+    with generation_cache_context():
+        # Build the API model object.
+        # This object is a frozen representation of the whole API, and is sent
+        # to each template in the rendering step.
+        api_schema = api.API.build(req.proto_file, opts=opts, package=package)
 
-    # Translate into a protobuf CodeGeneratorResponse; this reads the
-    # individual templates and renders them.
-    # If there are issues, error out appropriately.
-    res = generator.Generator(opts).get_response(api_schema, opts)
+        # Translate into a protobuf CodeGeneratorResponse; this reads the
+        # individual templates and renders them.
+        # If there are issues, error out appropriately.
+        res = generator.Generator(opts).get_response(api_schema, opts)
 
     # Output the serialized response.
     output.write(res.SerializeToString())

gapic/schema/api.py

@@ -115,7 +115,6 @@ 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] = None,
     ) -> "Proto":
         """Build and return a Proto instance.
 
@@ -140,7 +139,6 @@ def build(
             prior_protos=prior_protos or {},
             load_services=load_services,
             all_resources=all_resources or {},
-            context_cache=context_cache,
         ).proto
 
     @cached_property
@@ -456,7 +454,6 @@ def disambiguate_keyword_sanitize_fname(
         # type into the proto file that defines an LRO.
         # We just load all the APIs types first and then
         # load the services and methods with the full scope of types.
-        context_cache = {}
         pre_protos: Dict[str, Proto] = dict(prior_protos or {})
         for fd in file_descriptors:
             fd.name = disambiguate_keyword_sanitize_fname(fd.name, pre_protos)
@@ -468,7 +465,6 @@ 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.
@@ -489,7 +485,6 @@ 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()
         }
@@ -1108,7 +1103,6 @@ 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] = None,
     ):
         self.proto_messages: Dict[str, wrappers.MessageType] = {}
         self.proto_enums: Dict[str, wrappers.EnumType] = {}
@@ -1117,7 +1111,6 @@ 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.
         #
@@ -1227,28 +1220,20 @@ def proto(self) -> Proto:
         if not self.file_to_generate:
             return naive
 
-        global_collisions = frozenset(naive.names)
         visited_messages: Set[wrappers.MessageType] = set()
-        self.context_cache = {}
         # Return a context-aware proto object.
         return dataclasses.replace(
             naive,
             all_enums=collections.OrderedDict(
-                (
-                    k,
-                    v.with_context(
-                        collisions=global_collisions, context_cache=self.context_cache
-                    ),
-                )
+                (k, v.with_context(collisions=naive.names))
                 for k, v in naive.all_enums.items()
             ),
             all_messages=collections.OrderedDict(
                 (
                     k,
                     v.with_context(
-                        collisions=global_collisions,
+                        collisions=naive.names,
                         visited_messages=visited_messages,
-                        context_cache=self.context_cache,
                     ),
                 )
                 for k, v in naive.all_messages.items()
@@ -1261,14 +1246,11 @@ def proto(self) -> Proto:
                     v.with_context(
                         collisions=v.names,
                         visited_messages=visited_messages,
-                        context_cache=self.context_cache,
                     ),
                 )
                 for k, v in naive.services.items()
             ),
-            meta=naive.meta.with_context(
-                collisions=naive.names, context_cache=self.context_cache
-            ),
+            meta=naive.meta.with_context(collisions=naive.names),
         )
 
     @cached_property

gapic/schema/metadata.py

@@ -362,7 +362,7 @@ def resolve(self, selector: str) -> str:
 
     @cached_proto_context
     def with_context(
-        self, *, collisions: Set[str], context_cache: Optional[Dict] = None
+        self, *, collisions: Set[str]
     ) -> "Address":
         """Return a derivative of this address with the provided context.
 
@@ -404,7 +404,7 @@ def doc(self):
 
     @cached_proto_context
     def with_context(
-        self, *, collisions: Set[str], context_cache: Optional[Dict] = None
+        self, *, collisions: Set[str]
     ) -> "Metadata":
         """Return a derivative of this metadata with the provided context.
 
@@ -416,7 +416,7 @@ def with_context(
             dataclasses.replace(
                 self,
                 address=self.address.with_context(
-                    collisions=collisions, context_cache=context_cache
+                    collisions=collisions
                 ),
             )
             if collisions and collisions != self.address.collisions

gapic/schema/wrappers.py

@@ -417,7 +417,6 @@ def with_context(
         *,
         collisions: Set[str],
         visited_messages: Optional[Set["MessageType"]] = None,
-        context_cache: Optional[Dict] = None,
     ) -> "Field":
         """Return a derivative of this field with the provided context.
 
@@ -434,20 +433,19 @@ def with_context(
                         self.message in visited_messages if visited_messages else False
                     ),
                     visited_messages=visited_messages,
-                    context_cache=context_cache,
                 )
                 if self.message
                 else None
             ),
             enum=(
                 self.enum.with_context(
-                    collisions=collisions, context_cache=context_cache
+                    collisions=collisions
                 )
                 if self.enum
                 else None
             ),
             meta=self.meta.with_context(
-                collisions=collisions, context_cache=context_cache
+                collisions=collisions,
             ),
         )
 
@@ -751,7 +749,6 @@ def path_regex_str(self) -> str:
     def get_field(
         self,
         *field_path: str,
-        context_cache: Optional[Dict] = None,
         collisions: Optional[Set[str]] = None,
     ) -> Field:
         """Return a field arbitrarily deep in this message's structure.
@@ -796,7 +793,6 @@ def get_field(
             return cursor.with_context(
                 collisions=collisions,
                 visited_messages=set({self}),
-                context_cache=context_cache,
             )
 
         # Quick check: If cursor is a repeated field, then raise an exception.
@@ -828,7 +824,6 @@ def with_context(
         collisions: Set[str],
         skip_fields: bool = False,
         visited_messages: Optional[Set["MessageType"]] = None,
-        context_cache: Optional[Dict] = None,
     ) -> "MessageType":
         """Return a derivative of this message with the provided context.
 
@@ -849,28 +844,26 @@ def with_context(
                     k: v.with_context(
                         collisions=collisions,
                         visited_messages=visited_messages,
-                        context_cache=context_cache,
                     )
                     for k, v in self.fields.items()
                 }
                 if not skip_fields
                 else self.fields
             ),
             nested_enums={
-                k: v.with_context(collisions=collisions, context_cache=context_cache)
+                k: v.with_context(collisions=collisions)
                 for k, v in self.nested_enums.items()
             },
             nested_messages={
                 k: v.with_context(
                     collisions=collisions,
                     skip_fields=skip_fields,
                     visited_messages=visited_messages,
-                    context_cache=context_cache,
                 )
                 for k, v in self.nested_messages.items()
             },
             meta=self.meta.with_context(
-                collisions=collisions, context_cache=context_cache
+                collisions=collisions
             ),
         )
 
@@ -965,7 +958,6 @@ def with_context(
         self,
         *,
         collisions: Set[str],
-        context_cache: Optional[Dict] = None,
     ) -> "EnumType":
         """Return a derivative of this enum with the provided context.
 
@@ -977,7 +969,7 @@ def with_context(
             dataclasses.replace(
                 self,
                 meta=self.meta.with_context(
-                    collisions=collisions, context_cache=context_cache
+                    collisions=collisions
                 ),
             )
             if collisions
@@ -1095,7 +1087,6 @@ def with_context(
         *,
         collisions: Set[str],
         visited_messages: Optional[Set["MessageType"]] = None,
-        context_cache: Optional[Dict] = None,
     ) -> "ExtendedOperationInfo":
         """Return a derivative of this OperationInfo with the provided context.
 
@@ -1111,12 +1102,10 @@ def with_context(
                 request_type=self.request_type.with_context(
                     collisions=collisions,
                     visited_messages=visited_messages,
-                    context_cache=context_cache,
                 ),
                 operation_type=self.operation_type.with_context(
                     collisions=collisions,
                     visited_messages=visited_messages,
-                    context_cache=context_cache,
                 ),
             )
         )
@@ -1168,7 +1157,6 @@ def with_context(
         *,
         collisions: Set[str],
         visited_messages: Optional[Set["MessageType"]] = None,
-        context_cache: Optional[Dict] = None,
     ) -> "OperationInfo":
         """Return a derivative of this OperationInfo with the provided context.
 
@@ -1181,12 +1169,10 @@ def with_context(
             response_type=self.response_type.with_context(
                 collisions=collisions,
                 visited_messages=visited_messages,
-                context_cache=context_cache,
             ),
             metadata_type=self.metadata_type.with_context(
                 collisions=collisions,
                 visited_messages=visited_messages,
-                context_cache=context_cache,
             ),
         )
 
@@ -1982,7 +1968,6 @@ def with_context(
         *,
         collisions: Set[str],
         visited_messages: Optional[Set["MessageType"]] = None,
-        context_cache: Optional[Dict] = None,
     ) -> "Method":
         """Return a derivative of this method with the provided context.
 
@@ -1996,7 +1981,6 @@ def with_context(
                 self.lro.with_context(
                     collisions=collisions,
                     visited_messages=visited_messages,
-                    context_cache=context_cache,
                 )
                 if collisions
                 else self.lro
@@ -2006,7 +1990,6 @@ def with_context(
             self.extended_lro.with_context(
                 collisions=collisions,
                 visited_messages=visited_messages,
-                context_cache=context_cache,
             )
             if self.extended_lro
             else None
@@ -2019,15 +2002,13 @@ def with_context(
             input=self.input.with_context(
                 collisions=collisions,
                 visited_messages=visited_messages,
-                context_cache=context_cache,
             ),
             output=self.output.with_context(
                 collisions=collisions,
                 visited_messages=visited_messages,
-                context_cache=context_cache,
             ),
             meta=self.meta.with_context(
-                collisions=collisions, context_cache=context_cache
+                collisions=collisions,
             ),
         )
 
@@ -2410,7 +2391,6 @@ def with_context(
         *,
         collisions: Set[str],
         visited_messages: Optional[Set["MessageType"]] = None,
-        context_cache: Optional[Dict] = None,
     ) -> "Service":
         """Return a derivative of this service with the provided context.
 
@@ -2426,12 +2406,11 @@ def with_context(
                     # that may conflict with module imports.
                     collisions=collisions | set(v.flattened_fields.keys()),
                     visited_messages=visited_messages,
-                    context_cache=context_cache,
                 )
                 for k, v in self.methods.items()
             },
             meta=self.meta.with_context(
-                collisions=collisions, context_cache=context_cache
+                collisions=collisions,
             ),
         )