Replace pickle with JSON for cache serialization

Signed-off-by: Tamaki Nishino <otamachan@gmail.com>

Author: otamachan

rosidl_parser/rosidl_parser/cache.py

@@ -17,11 +17,75 @@
 import json
 import os
 import pathlib
-import pickle
 import shutil
 import tempfile
 from typing import Any, List, Optional, TypedDict
 
+import rosidl_parser.definition as _def
+
+# Build class registry and precomputed slots for safe deserialization
+_CLASSES = {}
+_SLOTS = {}
+
+for _name in dir(_def):
+    _obj = getattr(_def, _name)
+    if isinstance(_obj, type):
+        _CLASSES[_name] = _obj
+        # Precompute all slot names across MRO
+        _slots = []
+        for _klass in _obj.__mro__:
+            _klass_slots = getattr(_klass, '__slots__', ())
+            if isinstance(_klass_slots, str):
+                _slots.append(_klass_slots)
+            else:
+                _slots.extend(_klass_slots)
+        _SLOTS[_obj] = tuple(_slots)
+
+
+def _encode(obj):
+    if obj is None or isinstance(obj, (bool, int, float, str)):
+        return obj
+    if isinstance(obj, pathlib.Path):
+        return {'__path__': str(obj)}
+    if isinstance(obj, list):
+        return [_encode(item) for item in obj]
+    if isinstance(obj, tuple):
+        return {'__tuple__': [_encode(item) for item in obj]}
+    if isinstance(obj, dict):
+        return {k: _encode(v) for k, v in obj.items()}
+    slots = _SLOTS.get(type(obj))
+    if slots is None:
+        raise TypeError(f'Cannot JSON-encode object of type {type(obj).__name__}')
+    data = {'__class__': type(obj).__name__}
+    for slot in slots:
+        if hasattr(obj, slot):
+            data[slot] = _encode(getattr(obj, slot))
+    return data
+
+
+def _decode(data):
+    if data is None or isinstance(data, (bool, int, float, str)):
+        return data
+    if isinstance(data, list):
+        return [_decode(item) for item in data]
+    if isinstance(data, dict):
+        if '__path__' in data:
+            return pathlib.Path(data['__path__'])
+        if '__tuple__' in data:
+            return tuple(_decode(item) for item in data['__tuple__'])
+        cls_name = data.get('__class__')
+        if cls_name:
+            if cls_name not in _CLASSES:
+                raise ValueError(f'Unknown class: {cls_name}')
+            cls = _CLASSES[cls_name]
+            obj = cls.__new__(cls)
+            for slot in _SLOTS[cls]:
+                if slot in data:
+                    object.__setattr__(obj, slot, _decode(data[slot]))
+            return obj
+        return {k: _decode(v) for k, v in data.items()}
+    return data
+
 
 class CacheConfig(TypedDict):
     cache_dir: Optional[str]
@@ -157,7 +221,7 @@ def compute_cache_key(*args) -> Optional[str]:
         elif isinstance(arg, dict):
             hasher.update(json.dumps(arg, sort_keys=True).encode('utf-8'))
         else:
-            hasher.update(pickle.dumps(arg, protocol=pickle.HIGHEST_PROTOCOL))
+            hasher.update(json.dumps(arg, sort_keys=True, default=str).encode('utf-8'))
 
     return hasher.hexdigest()
 
@@ -171,9 +235,9 @@ def save_object_to_cache(cache_key: str, cache_subdir: str, obj: Any) -> None:
     tmp_dir = None
     try:
         tmp_dir = pathlib.Path(tempfile.mkdtemp(dir=cache_dir))
-        cache_file = tmp_dir / 'object.pkl'
-        with open(cache_file, 'wb') as f:
-            pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)
+        cache_file = tmp_dir / 'object.json'
+        with open(cache_file, 'w', encoding='utf-8') as f:
+            json.dump(_encode(obj), f)
         if cache_entry_dir.exists():
             shutil.rmtree(cache_entry_dir)
         tmp_dir.rename(cache_entry_dir)
@@ -190,12 +254,12 @@ def restore_object_from_cache(cache_key: str, cache_subdir: str) -> Optional[Any
     if not cache_dir:
         return None
     cache_entry_dir = cache_dir / cache_key
-    cache_file = cache_entry_dir / 'object.pkl'
+    cache_file = cache_entry_dir / 'object.json'
     if not cache_file.exists():
         return None
     try:
-        with open(cache_file, 'rb') as f:
-            return pickle.load(f)
+        with open(cache_file, 'r', encoding='utf-8') as f:
+            return _decode(json.load(f))
     except Exception as e:
         debug_print(f'[rosidl cache] Failed to load from cache: {e}')
         return None

rosidl_parser/test/test_cache.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import pathlib
 import time
 
 import pytest
@@ -136,3 +137,72 @@ def test_cleanup_removes_oldest(cache_dir):
 
     assert not old_entry.exists()
     assert new_entry.exists()
+
+
+def test_idl_file_cache_round_trip(cache_dir):
+    from rosidl_parser.definition import (
+        Array,
+        BasicType,
+        BoundedSequence,
+        BoundedString,
+        Constant,
+        IdlContent,
+        IdlFile,
+        IdlLocator,
+        Include,
+        Member,
+        Message,
+        NamespacedType,
+        Structure,
+        UnboundedSequence,
+        UnboundedString,
+    )
+
+    locator = IdlLocator(pathlib.Path('/base'), pathlib.Path('msg/Test.idl'))
+    content = IdlContent()
+
+    ns_type = NamespacedType(['test_pkg', 'msg'], 'TestMessage')
+    structure = Structure(ns_type, members=[
+        Member(BasicType('int32'), 'field_a'),
+        Member(BasicType('boolean'), 'field_b'),
+        Member(UnboundedString(), 'field_c'),
+        Member(BoundedString(100), 'field_d'),
+        Member(Array(BasicType('double'), 5), 'field_e'),
+        Member(BoundedSequence(BasicType('uint8'), 10), 'field_f'),
+        Member(UnboundedSequence(BasicType('int64')), 'field_g'),
+        Member(NamespacedType(['other_pkg', 'msg'], 'OtherType'), 'field_h'),
+    ])
+    msg = Message(structure)
+    msg.constants = [
+        Constant('MY_CONST', BasicType('int32'), 42),
+    ]
+
+    content.elements.append(Include('other_pkg/msg/OtherType.idl'))
+    content.elements.append(msg)
+
+    idl_file = IdlFile(locator, content)
+
+    save_object_to_cache('idl_key', 'sub', idl_file)
+    restored = restore_object_from_cache('idl_key', 'sub')
+
+    assert restored is not None
+    assert isinstance(restored, IdlFile)
+    assert str(restored.locator.basepath) == '/base'
+    assert str(restored.locator.relative_path) == 'msg/Test.idl'
+    assert len(restored.content.elements) == 2
+    assert isinstance(restored.content.elements[0], Include)
+    assert restored.content.elements[0].locator == 'other_pkg/msg/OtherType.idl'
+    msg_r = restored.content.elements[1]
+    assert isinstance(msg_r, Message)
+    assert msg_r.structure.namespaced_type.name == 'TestMessage'
+    assert len(msg_r.structure.members) == 8
+    assert isinstance(msg_r.structure.members[0].type, BasicType)
+    assert msg_r.structure.members[0].type.typename == 'int32'
+    assert isinstance(msg_r.structure.members[4].type, Array)
+    assert msg_r.structure.members[4].type.size == 5
+    assert isinstance(msg_r.structure.members[5].type, BoundedSequence)
+    assert msg_r.structure.members[5].type.maximum_size == 10
+    assert isinstance(msg_r.structure.members[6].type, UnboundedSequence)
+    assert len(msg_r.constants) == 1
+    assert msg_r.constants[0].name == 'MY_CONST'
+    assert msg_r.constants[0].value == 42