Python: Add Union/Intersection types, fix ref tracking, and address review items

- Add JavaType.Union and JavaType.Intersection with full sender/receiver/codec
  support (maps to JavaType$MultiCatch and JavaType$Intersection over RPC)
- Fix _send_as_ref to transfer ref tracking from before to after on CHANGE path
- Make _type_cache an instance variable to prevent cross-batch contamination
- Remove mkdir(parents=True) from _ensure_file_on_disk for safer path handling
- Reuse shutdown() in TyTypesClient.initialize for proper process lifecycle
- Add fully_qualified_name property to JavaType.Class and Parameterized
- Remove uv.lock from Git tracking

Author: knutwannheden

rewrite-python/rewrite/.gitignore

@@ -1,2 +1,3 @@
 .idea/
 .tox/
+uv.lock

rewrite-python/rewrite/src/rewrite/java/support_types.py

@@ -203,7 +203,7 @@ class Unknown(FullyQualified):
     class Class(FullyQualified):
         _flags_bit_map: int
         _fully_qualified_name: str
-        _kind: FullyQualified.Kind
+        _kind: JavaType.FullyQualified.Kind
         _type_parameters: Optional[List[JavaType]]
         _supertype: Optional[JavaType.FullyQualified]
         _owning_class: Optional[JavaType.FullyQualified]
@@ -212,6 +212,10 @@ class Class(FullyQualified):
         _members: Optional[List[JavaType.Variable]]
         _methods: Optional[List[JavaType.Method]]
 
+        @property
+        def fully_qualified_name(self) -> str:
+            return self._fully_qualified_name
+
     class ShallowClass(Class):
         pass
 
@@ -228,10 +232,10 @@ def type_parameters(self) -> Optional[List[JavaType]]:
             return self._type_parameters
 
         @property
-        def _fully_qualified_name(self) -> str:
+        def fully_qualified_name(self) -> str:
             t = getattr(self, '_type', None)
-            if t is not None and hasattr(t, '_fully_qualified_name'):
-                return t._fully_qualified_name
+            if t is not None and hasattr(t, 'fully_qualified_name'):
+                return t.fully_qualified_name
             return ''
 
     class GenericTypeVariable:
@@ -240,6 +244,24 @@ class Variance(Enum):
             Covariant = 1
             Contravariant = 2
 
+    @dataclass
+    class Union:
+        """Union type (e.g. str | int). Maps to JavaType$MultiCatch over RPC."""
+        _bounds: Optional[List[JavaType]] = field(default=None)
+
+        @property
+        def bounds(self) -> List[JavaType]:
+            return self._bounds if self._bounds is not None else []
+
+    @dataclass
+    class Intersection:
+        """Intersection type (e.g. A & B). Maps to JavaType$Intersection over RPC."""
+        _bounds: Optional[List[JavaType]] = field(default=None)
+
+        @property
+        def bounds(self) -> List[JavaType]:
+            return self._bounds if self._bounds is not None else []
+
     class Primitive(Enum):
         Boolean = 0
         Byte = 1

rewrite-python/rewrite/src/rewrite/java/support_types.pyi

@@ -77,6 +77,8 @@ class JavaType(ABC):
         _members: Optional[List[JavaType.Variable]]
         _methods: Optional[List[JavaType.Method]]
 
+        @property
+        def fully_qualified_name(self) -> str: ...
 
     class ShallowClass(Class):
         pass
@@ -85,12 +87,15 @@ class JavaType(ABC):
         _type: JavaType.FullyQualified
         _type_parameters: Optional[List[JavaType]]
 
+        @property
+        def fully_qualified_name(self) -> str: ...
+
         @property
         def type(self) -> JavaType.FullyQualified: ...
         @property
         def type_parameters(self) -> Optional[List[JavaType]]: ...
         @property
-        def _fully_qualified_name(self) -> str: ...
+        def fully_qualified_name(self) -> str: ...
 
 
     class GenericTypeVariable:
@@ -99,7 +104,17 @@ class JavaType(ABC):
             Covariant: Variance
             Contravariant: Variance
 
+    class Union:
+        _bounds: Optional[List[JavaType]]
 
+        @property
+        def bounds(self) -> List[JavaType]: ...
+
+    class Intersection:
+        _bounds: Optional[List[JavaType]]
+
+        @property
+        def bounds(self) -> List[JavaType]: ...
 
     class Primitive(Enum):
         Boolean: Primitive

rewrite-python/rewrite/src/rewrite/python/method_matcher.py

@@ -419,8 +419,6 @@ def _get_fqn(type_obj) -> Optional[str]:
     if type_obj is None:
         return None
 
-    if hasattr(type_obj, "_fully_qualified_name"):
-        return type_obj._fully_qualified_name
     if hasattr(type_obj, "fully_qualified_name"):
         return type_obj.fully_qualified_name
 

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

@@ -156,14 +156,7 @@ def initialize(self, project_root: str) -> bool:
             return True
 
         if self._initialized:
-            self._send_request("shutdown")
-            self._initialized = False
-            self._project_root = None
-            if self._process is not None:
-                try:
-                    self._process.wait(timeout=5)
-                except subprocess.TimeoutExpired:
-                    self._process.kill()
+            self.shutdown()
             self._start_process()
 
         result = self._send_request("initialize", {"projectRoot": project_root})

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

@@ -74,9 +74,6 @@ class PythonTypeMapping:
         method_type = mapping.method_invocation_type(call_node)
     """
 
-    # Cache for type mappings to avoid repeated class type creation
-    _type_cache: Dict[str, JavaType] = {}
-
     def __init__(self, source: str, file_path: Optional[str] = None, ty_client=None):
         """Initialize type mapping for a source file.
 
@@ -91,6 +88,7 @@ def __init__(self, source: str, file_path: Optional[str] = None, ty_client=None)
         self._file_path = file_path
         self._source_lines = source.splitlines()
         self._temp_file: Optional[Path] = None
+        self._type_cache: Dict[str, JavaType] = {}  # FQN -> JavaType (per-instance)
 
         # Compute line byte offsets for position conversion
         self._line_byte_offsets = self._compute_line_byte_offsets(source)
@@ -137,16 +135,18 @@ def _ensure_file_on_disk(self, source: str, file_path: Optional[str]) -> Optiona
         """Ensure the source is available as a file on disk for ty-types.
 
         Returns the absolute file path, or None if unavailable.
+        When file_path is given but doesn't exist, writes source there.
+        Callers are responsible for providing safe paths (e.g. within a temp directory).
         """
         if file_path:
             path = Path(file_path)
             if not path.is_absolute():
                 path = path.resolve()
             if path.exists():
                 return str(path)
-            # File path given but doesn't exist — write source there
+            # File path given but doesn't exist — write source there.
+            # The parent directory must already exist (caller should ensure this).
             try:
-                path.parent.mkdir(parents=True, exist_ok=True)
                 path.write_text(source, encoding='utf-8')
                 self._temp_file = path
                 return str(path)
@@ -177,7 +177,7 @@ def _compute_line_byte_offsets(source: str) -> List[int]:
     def _pos_to_byte_offset(self, lineno: int, col_offset: int) -> int:
         """Convert AST (lineno, col_offset) to an absolute byte offset.
 
-        Python's ast uses 1-based lineno and character-based col_offset.
+        Python's ast uses 1-based lineno and 0-based character col_offset (Python 3.8+).
         ty-types uses absolute byte offsets (ruff convention).
         """
         line_start = self._line_byte_offsets[lineno - 1]
@@ -276,7 +276,7 @@ def _resolve_type(self, type_id: int) -> Optional[JavaType]:
         if type_id in self._cycle_placeholders:
             placeholder = self._cycle_placeholders.pop(type_id)
             if isinstance(result, JavaType.Class):
-                placeholder._fully_qualified_name = result._fully_qualified_name
+                placeholder._fully_qualified_name = result.fully_qualified_name
                 placeholder._kind = result._kind
                 # Copy enriched fields so cycle placeholders retain supertypes/methods
                 for attr in ('_supertype', '_methods', '_type_parameters', '_interfaces',
@@ -285,8 +285,8 @@ def _resolve_type(self, type_id: int) -> Optional[JavaType]:
                     if val is not None:
                         setattr(placeholder, attr, val)
             elif isinstance(result, JavaType.Parameterized):
-                if hasattr(result._type, '_fully_qualified_name'):
-                    placeholder._fully_qualified_name = result._type._fully_qualified_name
+                if hasattr(result._type, 'fully_qualified_name'):
+                    placeholder._fully_qualified_name = result._type.fully_qualified_name
             self._type_id_cache[type_id] = placeholder
             return placeholder
 
@@ -351,16 +351,24 @@ def _descriptor_to_java_type(self, descriptor: Dict[str, Any]) -> Optional[JavaT
             return JavaType.Primitive.String
 
         elif kind == 'union':
-            # Unwrap union: take first non-None type
+            # Resolve all non-None members into a Union type.
+            # For Optional[X] (= X | None) with a single real member, unwrap to just X.
+            resolved_bounds = []
             for member_id in descriptor.get('members', []):
                 member = self._type_registry.get(member_id)
                 if member:
                     member_kind = member.get('kind')
                     # Skip None/NoneType members
                     if member_kind == 'instance' and member.get('className') in ('None', 'NoneType'):
                         continue
-                    return self._resolve_type(member_id)
-            return _UNKNOWN
+                    resolved = self._resolve_type(member_id)
+                    if resolved is not None:
+                        resolved_bounds.append(resolved)
+            if not resolved_bounds:
+                return _UNKNOWN
+            if len(resolved_bounds) == 1:
+                return resolved_bounds[0]
+            return JavaType.Union(_bounds=resolved_bounds)
 
         elif kind == 'module':
             module_name = descriptor.get('moduleName', '')
@@ -691,7 +699,12 @@ def _get_declaring_type(self, node: ast.Call) -> Optional[JavaType.FullyQualifie
         return self._infer_declaring_type_from_ast(node)
 
     def _resolve_declaring_type(self, type_id: int) -> Optional[JavaType.FullyQualified]:
-        """Resolve a type ID to a declaring type, maximizing object reuse."""
+        """Resolve a type ID to a declaring type, maximizing object reuse.
+
+        NOTE: The cycle-detection pattern here mirrors _resolve_type intentionally.
+        They use separate caches and placeholder dicts because declaring types are
+        resolved independently (often to a simpler Class without methods/members).
+        """
         if type_id in self._declaring_type_id_cache:
             return self._declaring_type_id_cache[type_id]
 
@@ -720,7 +733,7 @@ def _resolve_declaring_type(self, type_id: int) -> Optional[JavaType.FullyQualif
         if type_id in self._declaring_cycle_placeholders:
             placeholder = self._declaring_cycle_placeholders.pop(type_id)
             if isinstance(result, JavaType.Class):
-                placeholder._fully_qualified_name = result._fully_qualified_name
+                placeholder._fully_qualified_name = result.fully_qualified_name
                 placeholder._kind = result._kind
             self._declaring_type_id_cache[type_id] = placeholder
             return placeholder

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

@@ -921,6 +921,10 @@ def _receive_type(self, java_type, q: RpcReceiveQueue):
 
         This matches the sender's _visit_type which sends expanded type fields.
         The callback pattern ensures message counts match between sender and receiver.
+
+        IMPORTANT: isinstance ordering matters here. Parameterized must be checked
+        before Class because both are separate types. If inheritance changes, review
+        this ordering. Mirrors the sender's _visit_type chain.
         """
         from rewrite.java.support_types import JavaType as JT
 
@@ -994,7 +998,7 @@ def _receive_type(self, java_type, q: RpcReceiveQueue):
             #        owningClass, annotations, interfaces, members, methods
             flags = q.receive_defined(getattr(java_type, '_flags_bit_map', 0))
             kind = _to_enum(JT.FullyQualified.Kind)(q.receive(getattr(java_type, '_kind', JT.FullyQualified.Kind.Class)))
-            fqn = q.receive_defined(getattr(java_type, '_fully_qualified_name', ''))
+            fqn = q.receive_defined(getattr(java_type, 'fully_qualified_name', ''))
             type_params = q.receive_list(getattr(java_type, '_type_parameters', None) or [],
                                           lambda t: self._receive_type(t, q))
             supertype = q.receive(getattr(java_type, '_supertype', None),
@@ -1048,6 +1052,18 @@ def _receive_type(self, java_type, q: RpcReceiveQueue):
             var._annotations = annotations
             return var
 
+        elif isinstance(java_type, JT.Union):
+            # Union (MultiCatch in Java): bounds list
+            bounds = q.receive_list(getattr(java_type, '_bounds', None) or [],
+                                     lambda t: self._receive_type(t, q))
+            return JT.Union(_bounds=bounds)
+
+        elif isinstance(java_type, JT.Intersection):
+            # Intersection: bounds list
+            bounds = q.receive_list(getattr(java_type, '_bounds', None) or [],
+                                     lambda t: self._receive_type(t, q))
+            return JT.Intersection(_bounds=bounds)
+
         elif isinstance(java_type, JT.Unknown):
             # Unknown has no additional fields
             return java_type
@@ -1361,7 +1377,7 @@ def _receive_java_type_class(cls, q: RpcReceiveQueue):
     # owningClass, annotations, interfaces, members, methods
     flags = q.receive_defined(getattr(cls, '_flags_bit_map', 0) if cls else 0)
     kind = _to_enum(JT.FullyQualified.Kind)(q.receive(getattr(cls, '_kind', JT.FullyQualified.Kind.Class) if cls else JT.FullyQualified.Kind.Class))
-    fqn = q.receive_defined(getattr(cls, '_fully_qualified_name', '') if cls else '')
+    fqn = q.receive_defined(getattr(cls, 'fully_qualified_name', '') if cls else '')
     type_params = q.receive_list(getattr(cls, '_type_parameters', None) if cls else None)
     supertype = q.receive(getattr(cls, '_supertype', None) if cls else None)
     owning_class = q.receive(getattr(cls, '_owning_class', None) if cls else None)
@@ -1431,6 +1447,22 @@ def _receive_java_type_variable(variable, q: RpcReceiveQueue):
     return var
 
 
+def _receive_java_type_union(union, q: RpcReceiveQueue):
+    """Codec for receiving JavaType.Union (MultiCatch in Java) - consumes bounds list."""
+    from rewrite.java.support_types import JavaType as JT
+
+    bounds = q.receive_list(union._bounds)
+    return JT.Union(_bounds=bounds)
+
+
+def _receive_java_type_intersection(intersection, q: RpcReceiveQueue):
+    """Codec for receiving JavaType.Intersection - consumes bounds list."""
+    from rewrite.java.support_types import JavaType as JT
+
+    bounds = q.receive_list(intersection._bounds)
+    return JT.Intersection(_bounds=bounds)
+
+
 def _register_java_type_codecs():
     """Register codecs for JavaType classes."""
     from rewrite.java.support_types import JavaType as JT
@@ -1493,6 +1525,22 @@ def _register_java_type_codecs():
         lambda: JT.Array()  # Factory creates empty Array
     )
 
+    # JavaType.Union (MultiCatch in Java) - bounds list
+    register_codec_with_both_names(
+        'org.openrewrite.java.tree.JavaType$MultiCatch',
+        JT.Union,
+        _receive_java_type_union,
+        lambda: JT.Union()  # Factory creates empty Union
+    )
+
+    # JavaType.Intersection - bounds list
+    register_codec_with_both_names(
+        'org.openrewrite.java.tree.JavaType$Intersection',
+        JT.Intersection,
+        _receive_java_type_intersection,
+        lambda: JT.Intersection()  # Factory creates empty Intersection
+    )
+
 
 def _register_tree_codecs():
     """Register codecs for all AST types using reflection."""

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

@@ -875,7 +875,7 @@ def _visit_type(self, java_type, q: 'RpcSendQueue') -> None:
             #        owningClass, annotations, interfaces, members, methods
             q.get_and_send(java_type, lambda x: getattr(x, '_flags_bit_map', 0))
             q.get_and_send(java_type, lambda x: getattr(x, '_kind', JT.FullyQualified.Kind.Class))
-            q.get_and_send(java_type, lambda x: getattr(x, '_fully_qualified_name', ''))
+            q.get_and_send(java_type, lambda x: getattr(x, 'fully_qualified_name', ''))
             q.get_and_send_list_as_ref(java_type, lambda x: getattr(x, '_type_parameters', None) or [], self._type_signature, lambda t: self._visit_type(t, q))
             q.get_and_send_as_ref(java_type, lambda x: getattr(x, '_supertype', None), lambda t: self._visit_type(t, q))
             q.get_and_send_as_ref(java_type, lambda x: getattr(x, '_owning_class', None), lambda t: self._visit_type(t, q))
@@ -896,6 +896,14 @@ def _visit_type(self, java_type, q: 'RpcSendQueue') -> None:
             q.get_and_send_as_ref(java_type, lambda x: x._type, lambda t: self._visit_type(t, q))
             q.get_and_send_list_as_ref(java_type, lambda x: x._annotations or [], self._type_signature, lambda t: self._visit_type(t, q))
 
+        elif isinstance(java_type, JT.Union):
+            # Union (MultiCatch in Java): bounds list
+            q.get_and_send_list_as_ref(java_type, lambda x: x.bounds, self._type_signature, lambda t: self._visit_type(t, q))
+
+        elif isinstance(java_type, JT.Intersection):
+            # Intersection: bounds list
+            q.get_and_send_list_as_ref(java_type, lambda x: x.bounds, self._type_signature, lambda t: self._visit_type(t, q))
+
         elif isinstance(java_type, JT.Unknown):
             # Unknown has no additional fields
             pass
@@ -909,7 +917,7 @@ def _type_signature(self, java_type) -> str:
         if isinstance(java_type, JT.Primitive):
             return java_type.name
         if isinstance(java_type, JT.Class):
-            return getattr(java_type, '_fully_qualified_name', str(id(java_type)))
+            return getattr(java_type, 'fully_qualified_name', str(id(java_type)))
         if isinstance(java_type, JT.Method):
             declaring = getattr(java_type, '_declaring_type', None)
             declaring_name = self._type_signature(declaring) if declaring else ''
@@ -924,6 +932,10 @@ def _type_signature(self, java_type) -> str:
         if isinstance(java_type, JT.Array):
             elem_sig = self._type_signature(java_type._elem_type) if java_type._elem_type else ''
             return f"{elem_sig}[]"
+        if isinstance(java_type, JT.Union):
+            return '|'.join(self._type_signature(b) for b in java_type.bounds)
+        if isinstance(java_type, JT.Intersection):
+            return '&'.join(self._type_signature(b) for b in java_type.bounds)
         return str(id(java_type))
 
     def _visit_space(self, space: Space, q: 'RpcSendQueue') -> None: