chore: cache proto context to optimize generator performance #2531

ohmayr · 2026-01-21T23:25:54Z

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
Ensure the tests and linter pass
Code coverage does not decrease (if any source code was changed)
Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

gemini-code-assist · 2026-01-21T23:26:11Z

Summary of Changes

Hello @ohmayr, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request implements a significant performance optimization by introducing a caching mechanism for context-aware schema objects. By passing a shared context_cache through the with_context methods, the system can avoid recomputing identical object derivatives, especially those involved in resolving naming collisions. This change aims to reduce processing time and improve the overall efficiency of the GAPIC generator.

Highlights

Introduced Context Caching: A context_cache dictionary is now passed through various with_context methods across api.py, metadata.py, and wrappers.py to store and reuse previously computed context-aware objects, preventing redundant calculations.
Optimized with_context Calls: The with_context methods in Address, Metadata, Field, MessageType, EnumType, ExtendedOperationInfo, OperationInfo, Method, and Service now check the context_cache before performing computations, returning cached results if available.
Early Exit for MessageType: An optimization was added to MessageType.with_context to return the current instance early if the collisions set has not changed, further reducing redundant processing.
Global Collisions as Frozenset: In api.py, naive.names is now converted to a frozenset (global_collisions) once and reused, improving efficiency when passed to multiple with_context calls.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a caching mechanism to optimize with_context calls, which is a good idea for performance. However, the implementation has a few issues that need to be addressed: a critical bug in api.py that will cause a runtime error, and several issues with the caching implementation itself regarding type safety, potential cache corruption, and code duplication. My detailed comments point out these issues with suggestions for fixes.

gapic/schema/api.py

gapic/schema/metadata.py

daniel-sanche · 2026-01-23T21:11:51Z

gapic/schema/api.py


+        global_collisions = frozenset(naive.names)
        visited_messages: Set[wrappers.MessageType] = set()
+        self.context_cache = {}


Why is this state being reset here?

Is the dict passed in at init time used anywhere? This seems to be the only usage?

daniel-sanche · 2026-01-23T21:18:14Z

gapic/utils/cache.py

+    """Decorator to memoize with_context calls based on self and collisions."""
+
+    @functools.wraps(func)
+    def wrapper(self, *, collisions, context_cache: Optional[Dict] = None, **kwargs):


How does this interact with the visited_messages argument?

I'm having trouble understanding what visited_messages is used for, but this cache doesn't seem to take it into account. Is that ok?

daniel-sanche · 2026-01-23T21:31:58Z

gapic/utils/cache.py

+
+        # 2. Create the cache key
+        collisions_key = frozenset(collisions) if collisions else None
+        key = (id(self), collisions_key)


With this key, the cache will only be triggered if both messages have the same memory address, even if the content is identical. Is that the intention? Typically, a hash would be used for this kind of thing. Unless each message is a singleton

a = ("hello",) b = ("hello",) >>> id(a), id(b) (140612776488128, 140612776488000) >>> hash(a), hash(b) (2145482566216562249, 2145482566216562249)

daniel-sanche · 2026-01-23T21:39:37Z

gapic/utils/cache.py

+    def wrapper(self, *, collisions, context_cache: Optional[Dict] = None, **kwargs):
+        # 1. Initialize cache if not provided (handles the root call case)
+        if context_cache is None:
+            context_cache = {}


Does context_cache need to be optional? It seems to me that the root call creates a new dictionary, if I'm reading it right

product-auto-label bot added the size: m Pull request size is medium. label Jan 21, 2026

gemini-code-assist bot reviewed Jan 21, 2026

View reviewed changes

gapic/schema/api.py Outdated Show resolved Hide resolved

gapic/schema/metadata.py Outdated Show resolved Hide resolved

gapic/schema/metadata.py Outdated Show resolved Hide resolved

gapic/schema/metadata.py Outdated Show resolved Hide resolved

ohmayr force-pushed the cache-api-context branch from 68e453e to 8e7e903 Compare January 22, 2026 01:36

ohmayr added 6 commits January 22, 2026 23:33

optimize gapic

9fa9487

update type to Dict

afaa050

reduce diff

edad384

use module level cache

2103a8a

pass cache to meta

fa92e02

fix positional argument

58f3049

ohmayr force-pushed the cache-api-context branch from ddea9d8 to 58f3049 Compare January 22, 2026 23:35

pass down service names

4965011

ohmayr force-pushed the cache-api-context branch from cc9e070 to 4965011 Compare January 23, 2026 09:14

fix lint

e03a978

daniel-sanche reviewed Jan 23, 2026

View reviewed changes

avoid resetting the context cache

fd827c6

ohmayr changed the title ~~optimize gapic~~ chore: cache proto context to optimize generator performance Jan 23, 2026

ohmayr added 2 commits January 23, 2026 22:40

define cache in API layer

fd5c07b

fix init file

2537d98

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

chore: cache proto context to optimize generator performance #2531

chore: cache proto context to optimize generator performance #2531

ohmayr commented Jan 21, 2026

Uh oh!

gemini-code-assist bot commented Jan 21, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

daniel-sanche Jan 23, 2026

Uh oh!

daniel-sanche Jan 23, 2026

Uh oh!

daniel-sanche Jan 23, 2026

Uh oh!

daniel-sanche Jan 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

chore: cache proto context to optimize generator performance #2531

Are you sure you want to change the base?

chore: cache proto context to optimize generator performance #2531

Conversation

ohmayr commented Jan 21, 2026

Uh oh!

gemini-code-assist bot commented Jan 21, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

daniel-sanche Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

daniel-sanche Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

daniel-sanche Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

daniel-sanche Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants