Intern JavaType instances across a parse session and keep RPC refs across files

rewrite-python/rewrite/src/rewrite/python/ty_client.py

@@ -34,6 +34,8 @@
 from pathlib import Path
 from typing import Any, Dict, Optional
 
+from .type_mapping import SessionTypeCache
+
 logger = logging.getLogger(__name__)
 
 
@@ -92,6 +94,9 @@ def __init__(self, virtual_env: Optional[str] = None):
         # state leaks across unrelated parses.
         self.session_types: Dict[int, Dict[str, Any]] = {}
 
+        # The JavaTypes those descriptors resolve to, sharing their lifetime.
+        self.java_types = SessionTypeCache()
+
         self._start_process()
 
     def __enter__(self) -> TyTypesClient:
@@ -294,6 +299,7 @@ def initialize(self, project_root: str) -> bool:
             # A different project root means a brand-new ty session whose type
             # ids start over; drop the accumulated table so ids don't collide.
             self.session_types.clear()
+            self.java_types.clear()
             self._start_process()
 
         # Bounded so a wedged ty degrades to untyped instead of hanging; large

rewrite-python/rewrite/src/rewrite/python/type_mapping.py

@@ -27,6 +27,7 @@
 import ast
 import os
 import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Tuple
 
@@ -137,6 +138,25 @@ def _is_public(name: str, all_names: Optional[Set[str]]) -> bool:
     return name in all_names if all_names is not None else not name.startswith('_')
 
 
+@dataclass
+class SessionTypeCache:
+    """The JavaType instances every file of one ty session shares.
+
+    ty type ids are stable for a session's lifetime and its descriptor table is
+    cumulative across files (see ``TyTypesClient.session_types``), so an id denotes
+    the same type in every file of the parse and needs only one JavaType.
+    """
+
+    by_type_id: Dict[int, JavaType] = field(default_factory=dict)
+    declaring_by_type_id: Dict[int, JavaType.FullyQualified] = field(default_factory=dict)
+    by_fqn: Dict[str, JavaType] = field(default_factory=dict)
+
+    def clear(self) -> None:
+        self.by_type_id.clear()
+        self.declaring_by_type_id.clear()
+        self.by_fqn.clear()
+
+
 class PythonTypeMapping:
     """Maps Python types to JavaType for recipe matching.
 
@@ -170,7 +190,6 @@ def __init__(self, source: str, file_path: Optional[str] = None, ty_client=None,
         self._source = source
         self._file_path = file_path
         self._temp_file: Optional[Path] = None
-        self._type_cache: Dict[str, JavaType] = {}  # FQN -> JavaType (per-instance)
 
         # Use pre-computed values when available (e.g. supplied by ParserVisitor),
         # otherwise compute them here.
@@ -189,8 +208,13 @@ def __init__(self, source: str, file_path: Optional[str] = None, ty_client=None,
         self._node_index_by_start: Dict[int, List[Tuple[int, int, str]]] = {}  # start -> [(end, type_id, node_kind)]
         self._type_registry: Dict[int, Dict[str, Any]] = {}  # type_id -> TypeDescriptor
         self._call_signature_index: Dict[Tuple[int, int], Dict[str, Any]] = {}  # (start, end) -> callSignature
-        self._type_id_cache: Dict[int, JavaType] = {}  # type_id -> resolved JavaType
-        self._declaring_type_id_cache: Dict[int, JavaType.FullyQualified] = {}  # type_id -> resolved declaring type
+        session_java_types = getattr(ty_client, 'java_types', None) or SessionTypeCache()
+        self._type_cache: Dict[str, JavaType] = session_java_types.by_fqn
+        self._type_id_cache: Dict[int, JavaType] = session_java_types.by_type_id
+        self._declaring_type_id_cache: Dict[int, JavaType.FullyQualified] = \
+            session_java_types.declaring_by_type_id
+        # Cycle detection tracks one file's in-progress resolutions, so unlike the
+        # resolved types above it belongs to this instance rather than the session.
         self._resolving_type_ids: set = set()  # type_ids currently being resolved (cycle detection)
         self._resolving_declaring_type_ids: set = set()
         self._cycle_placeholders: Dict[int, JavaType.Class] = {}  # placeholders created on cycle detection
@@ -643,9 +667,8 @@ def _descriptor_to_java_type(self, descriptor: Dict[str, Any]) -> Optional[JavaT
             # Map a TypedDict to a nominal class type by name and populate its
             # members from the descriptor's `fields`. Each field carries its own
             # `name` and `typeId` (the same shape as a classLiteral member), so
-            # we reuse the variable-building path. The class is keyed by simple
-            # name via `_create_class_type`, so two TypedDicts that share a name
-            # collapse — acceptable until ty emits a qualified name here.
+            # we reuse the variable-building path. Two TypedDicts sharing a name
+            # within a file collapse (see _typed_dict_key).
             #
             # We still drop the PEP 728 `closed` / `extraItems` openness fields
             # and the per-field `required` / `readOnly` flags; and linking a
@@ -655,7 +678,8 @@ def _descriptor_to_java_type(self, descriptor: Dict[str, Any]) -> Optional[JavaT
             name = descriptor.get('name', '')
             if not name:
                 return _UNKNOWN
-            class_type = self._create_class_type(name, shallow=False)
+            class_type = self._create_class_type(
+                name, shallow=False, cache_key=self._typed_dict_key(name))
             fields = descriptor.get('fields', [])
             if fields and getattr(class_type, '_members', None) is None:
                 variables = []
@@ -1337,7 +1361,7 @@ def _declaring_type_from_descriptor(self, descriptor: Dict[str, Any]) -> Optiona
         elif kind == 'typedDict':
             name = descriptor.get('name', '')
             if name:
-                return self._create_class_type(name)
+                return self._create_class_type(name, cache_key=self._typed_dict_key(name))
             return None
 
         elif kind == 'subclassOf':
@@ -1569,11 +1593,19 @@ def _make_class_object_type(self, base: JavaType) -> JavaType:
         param._type_parameters = [base]
         return param
 
-    def _create_class_type(self, fqn: str, shallow: bool = True) -> JavaType.Class:
+    def _typed_dict_key(self, name: str) -> str:
+        """The cache key for a TypedDict named ``name``. ty names a TypedDict
+        without qualifying it, so the key is scoped to this file to keep the same
+        name in two modules apart in the session-wide cache."""
+        return f"{self._file_path}#{name}"
+
+    def _create_class_type(self, fqn: str, shallow: bool = True,
+                           cache_key: Optional[str] = None) -> JavaType.Class:
         """Create a class type from a fully qualified name. Stubs minted here are
         body-less references (ShallowClass); pass shallow=False when filling a body."""
-        if fqn in self._type_cache:
-            cached = self._type_cache[fqn]
+        key = fqn if cache_key is None else cache_key
+        if key in self._type_cache:
+            cached = self._type_cache[key]
             if isinstance(cached, JavaType.Class):
                 if not shallow and type(cached) is JavaType.ShallowClass:
                     # Promote in place so earlier references see the full class.
@@ -1585,7 +1617,7 @@ def _create_class_type(self, fqn: str, shallow: bool = True) -> JavaType.Class:
         class_type._fully_qualified_name = fqn
         class_type._kind = JavaType.FullyQualified.Kind.Class
 
-        self._type_cache[fqn] = class_type
+        self._type_cache[key] = class_type
         return class_type
 
     def _get_node_text(self, node: ast.expr) -> str:

rewrite-python/rewrite/src/rewrite/rpc/reference.py

@@ -0,0 +1,63 @@
+# Copyright 2025 the original author or authors.
+# <p>
+# Licensed under the Moderne Source Available License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# <p>
+# https://docs.moderne.io/licensing/moderne-source-available-license
+# <p>
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Ref-id bookkeeping for the send side of the RPC protocol."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple
+
+
+class ReferenceMap:
+    """Ref ids assigned to objects sent to a peer, keyed by object identity.
+
+    The peer builds the mirror of this map as it receives, so an id may only be
+    cited once the object behind it has arrived. Its lifetime is therefore the
+    peer connection's, and ``rollback_to`` undoes the ids of an exchange the peer
+    never received. Mirrors ``RewriteRpc.localRefs`` and the JavaScript and Go
+    ``ReferenceMap``.
+    """
+
+    def __init__(self) -> None:
+        # The object is kept beside its ref so a recycled id() cannot alias a
+        # freed object onto another object's ref.
+        self._by_id: Dict[int, Tuple[Any, int]] = {}
+        self._next: int = 0
+
+    def get(self, obj: Any) -> Optional[int]:
+        """The ref already assigned to ``obj``, or None when it has not been sent."""
+        entry = self._by_id.get(id(obj))
+        return entry[1] if entry is not None and entry[0] is obj else None
+
+    def create(self, obj: Any) -> int:
+        self._next += 1
+        self._by_id[id(obj)] = (obj, self._next)
+        return self._next
+
+    def snapshot(self) -> int:
+        """The high-water ref id, for a later ``rollback_to``."""
+        return self._next
+
+    def rollback_to(self, snapshot: int) -> None:
+        """Drop every ref assigned since ``snapshot``."""
+        for key in [k for k, (_, ref) in self._by_id.items() if ref > snapshot]:
+            del self._by_id[key]
+        self._next = snapshot
+
+    def clear(self) -> None:
+        self._by_id.clear()
+        self._next = 0
+
+    def __len__(self) -> int:
+        return len(self._by_id)

rewrite-python/rewrite/src/rewrite/rpc/send_queue.py

@@ -13,6 +13,8 @@
 from typing import Any, Dict, List, Optional, Callable, TypeVar
 from uuid import UUID
 
+from rewrite.rpc.reference import ReferenceMap
+
 
 class RpcObjectState(str, Enum):
     NO_CHANGE = "NO_CHANGE"
@@ -30,10 +32,11 @@ class RpcObjectState(str, Enum):
 class RpcSendQueue:
     """Queue for generating RpcObjectData array from Python LST using visitor pattern."""
 
-    def __init__(self, source_file_type: Optional[str] = None):
+    def __init__(self, source_file_type: Optional[str] = None,
+                 refs: Optional[ReferenceMap] = None):
         self.q: List[Dict[str, Any]] = []
-        self.refs: Dict[int, tuple] = {}  # id(obj) -> (obj, ref_number) — verified with `is`
-        self.next_ref: int = 0
+        # A caller-supplied map spans the peer connection; the default spans this queue.
+        self.refs: ReferenceMap = refs if refs is not None else ReferenceMap()
         self.source_file_type = source_file_type
         self._before: Any = None
 
@@ -220,21 +223,19 @@ def _add_as_ref(self, obj: Any, on_change: Optional[Callable[[], None]] = None)
             self.put({'state': RpcObjectState.DELETE})
             return
 
-        obj_id = id(obj)
-        entry = self.refs.get(obj_id)
-        if entry is not None and entry[0] is obj:
+        ref = self.refs.get(obj)
+        if ref is not None:
             # Already sent — emit ref number only, no onChange
-            self.put({'state': RpcObjectState.ADD, 'ref': entry[1]})
+            self.put({'state': RpcObjectState.ADD, 'ref': ref})
             return
 
         # First time — assign ref number and serialize fully
-        self.next_ref += 1
-        self.refs[obj_id] = (obj, self.next_ref)
+        ref = self.refs.create(obj)
 
         value_type = self._get_value_type(obj)
         codec = self._get_rpc_codec(obj)
         value = None if on_change is not None or codec is not None else self._get_primitive_value(obj)
-        self.put({'state': RpcObjectState.ADD, 'valueType': value_type, 'value': value, 'ref': self.next_ref})
+        self.put({'state': RpcObjectState.ADD, 'valueType': value_type, 'value': value, 'ref': ref})
         self._do_change(obj, None, on_change, codec)
 
     def _do_change(self, after: Any, before: Any,