Refactor serialization: Enhance type safety by introducing Optional in various data structures, implement LanguageVersion interning for memory optimization, and update Protocol Buffer handling for improved deserialization logic.

Author: ftomassetti

README.md

@@ -65,5 +65,5 @@ coverage html
 ## Update Protobuffer classes
 
 ```
-protoc --proto_path=./src --python_out=./src ./src/lionweb/serialization/proto/Chunk.proto
+protoc --proto_path=./src --python_out=./src --mypy_out=./src -I . ./src/lionweb/serialization/proto/Chunk.proto 
 ```

requirements.txt

@@ -45,6 +45,7 @@ mdurl==0.1.2
 more-itertools==10.6.0
 mypy==1.15.0
 mypy-extensions==1.0.0
+mypy-protobuf==3.6.0
 nh3==0.2.20
 nodeenv==1.9.1
 numpy==2.0.2

src/lionweb/client/repository_archives.py

@@ -42,5 +42,3 @@ def upload(bulk_import: BulkImport) -> int:
     end = time.perf_counter()
     elapsed_seconds = end - start
     print(f"Uploaded {total_nodes} nodes in {elapsed_seconds:.3f} seconds")
-
-

src/lionweb/serialization/archive.py

@@ -9,7 +9,7 @@
 def load_archive(filename) -> List[SerializationChunk]:
     ps = ProtoBufSerialization(LionWebVersion.V2023_1)
     start = time.time()
-    chunks = []
+    chunks : List[SerializationChunk] = []
     import zipfile
     with zipfile.ZipFile(filename, 'r') as zf:
         for name in zf.namelist():

src/lionweb/serialization/data/language_version.py

@@ -1,15 +1,45 @@
-from typing import Optional
+from typing import Optional, Dict, Tuple
+import threading
 
 
 class LanguageVersion:
     """
     The pair Language Key and Language Version identify a specific version of a language.
-    Corresponds to the Java class 'UsedLanguage'.
+    It is used also in the role of 'UsedLanguage', as specified in the specs.
     """
 
+    # Class-level cache for interning instances
+    _instances: Dict[Tuple[Optional[str], Optional[str]], 'LanguageVersion'] = {}
+    _lock = threading.Lock()  # Thread-safe access to cache
+
+    def __new__(cls, key: Optional[str] = None, version: Optional[str] = None):
+        # Create cache key
+        cache_key = (key, version)
+        
+        # Thread-safe cache lookup
+        with cls._lock:
+            if cache_key in cls._instances:
+                return cls._instances[cache_key]
+            
+            # Create new instance and cache it
+            instance = super().__new__(cls)
+            cls._instances[cache_key] = instance
+            return instance
+
     def __init__(self, key: Optional[str] = None, version: Optional[str] = None):
-        self.key = key
-        self.version = version
+        # Only initialize if not already initialized (due to interning)
+        if not hasattr(self, '_initialized'):
+            self._key = key
+            self._version = version
+            self._initialized = True
+
+    @classmethod
+    def of(cls, key: Optional[str] = None, version: Optional[str] = None) -> 'LanguageVersion':
+        """
+        Factory method to get an interned LanguageVersion instance.
+        This is the preferred way to create LanguageVersion instances.
+        """
+        return cls(key, version)
 
     @staticmethod
     def from_language(language):
@@ -51,23 +81,46 @@ def from_meta_pointer(meta_pointer):
             raise ValueError("meta_pointer language should not be null")
         if meta_pointer.version is None:
             raise ValueError("meta_pointer version should not be null")
-        return LanguageVersion(meta_pointer.language, meta_pointer.version)
+        return LanguageVersion.of(meta_pointer.language, meta_pointer.version)
 
     def get_key(self) -> Optional[str]:
-        return self.key
+        return self._key
 
     def set_key(self, key: str):
-        self.key = key
+        raise RuntimeError("LanguageVersion instances are immutable after creation")
 
     def get_version(self) -> Optional[str]:
-        return self.version
+        return self._version
 
     def set_version(self, version: str):
-        self.version = version
+        raise RuntimeError("LanguageVersion instances are immutable after creation")
+
+    @property
+    def key(self) -> Optional[str]:
+        return self._key
+
+    @property
+    def version(self) -> Optional[str]:
+        return self._version
+
+    @classmethod
+    def clear_cache(cls):
+        """Clear the interning cache. Useful for testing or memory management."""
+        with cls._lock:
+            cls._instances.clear()
+
+    @classmethod
+    def cache_size(cls) -> int:
+        """Get the current size of the interning cache."""
+        with cls._lock:
+            return len(cls._instances)
 
     def __eq__(self, other):
         if not isinstance(other, LanguageVersion):
             return False
+        # With interning, we can use identity comparison for performance
+        if self is other:
+            return True
         return self.key == other.key and self.version == other.version
 
     def __hash__(self):

src/lionweb/serialization/data/serialized_chunk.py

@@ -27,7 +27,7 @@ def get_instance_by_id(self, instance_id: str) -> SerializedClassifierInstance:
             raise ValueError(f"Cannot find instance with ID {instance_id}")
         return instance
 
-    def add_language(self, language):
+    def add_language(self, language: LanguageVersion) -> None:
         self.languages.append(language)
 
     def __str__(self):

src/lionweb/serialization/data/serialized_classifier_instance.py

@@ -17,7 +17,7 @@ class SerializedClassifierInstance:
     properties: List[SerializedPropertyValue] = field(default_factory=list)
     containments: List[SerializedContainmentValue] = field(default_factory=list)
     references: List[SerializedReferenceValue] = field(default_factory=list)
-    annotations: List[str] = field(default_factory=list)
+    annotations: List[Optional[str]] = field(default_factory=list)
     parent_node_id: Optional[str] = None
 
     def get_parent_node_id(self):
@@ -57,7 +57,7 @@ def set_property_value(self, property_meta_pointer, serialized_value):
             SerializedPropertyValue(property_meta_pointer, serialized_value)
         )
 
-    def add_children(self, containment_meta_pointer, children_ids: List[str]):
+    def add_children(self, containment_meta_pointer, children_ids: List[Optional[str]]):
         from .serialized_containment_value import SerializedContainmentValue
 
         self.containments.append(
@@ -101,22 +101,22 @@ def get_reference_values(self, reference_meta_pointer) -> List:
                 return rv.get_value()
         return []
 
-    def get_containment_values_by_key(self, containment_key: str) -> List[str]:
+    def get_containment_values_by_key(self, containment_key: str) -> List[Optional[str]]:
         for rv in self.containments:
             if rv.get_meta_pointer().key == containment_key:
                 return rv.get_value()
         return []
 
-    def get_containment_values(self, containment_meta_pointer) -> List[str]:
+    def get_containment_values(self, containment_meta_pointer) -> List[Optional[str]]:
         for cv in self.containments:
             if containment_meta_pointer == cv.get_meta_pointer():
                 return cv.get_value()
         return []
 
-    def set_annotations(self, annotation_ids: List[str]):
+    def set_annotations(self, annotation_ids: List[Optional[str]]):
         self.annotations = annotation_ids[:]
 
-    def add_annotation(self, annotation_id: str):
+    def add_annotation(self, annotation_id: Optional[str]):
         self.annotations.append(annotation_id)
 
     def __eq__(self, other):

src/lionweb/serialization/data/serialized_containment_value.py

@@ -1,10 +1,10 @@
-from typing import List
+from typing import List, Optional
 
 from lionweb.serialization.data.metapointer import MetaPointer
 
 
 class SerializedContainmentValue:
-    def __init__(self, meta_pointer: MetaPointer, value: List[str]):
+    def __init__(self, meta_pointer: MetaPointer, value: List[Optional[str]]):
         self.meta_pointer = meta_pointer
         self.value = value if value is not None else []
 
@@ -14,10 +14,10 @@ def get_meta_pointer(self) -> MetaPointer:
     def set_meta_pointer(self, meta_pointer):
         self.meta_pointer = meta_pointer
 
-    def get_value(self) -> List[str]:
+    def get_value(self) -> List[Optional[str]]:
         return self.value.copy()
 
-    def set_value(self, value: List[str]):
+    def set_value(self, value: List[Optional[str]]):
         self.value = value.copy()
 
     def __eq__(self, other):

src/lionweb/serialization/low_level_json_serialization.py

@@ -1,5 +1,5 @@
 import json
-from typing import Iterable, List, cast
+from typing import Iterable, List, cast, Optional
 
 from lionweb import LionWebVersion
 from lionweb.serialization.data.metapointer import MetaPointer
@@ -190,7 +190,6 @@ def _read_languages(
         if isinstance(languages, list):
             for element in languages:
                 try:
-                    language_key_version = LanguageVersion()
                     if isinstance(element, dict):
                         extra_keys = set(element.keys()) - {"key", "version"}
                         if extra_keys:
@@ -207,8 +206,7 @@ def _read_languages(
                             raise ValueError(
                                 "Both 'key' and 'version' should be strings"
                             )
-                        language_key_version.key = element.get("key")
-                        language_key_version.version = element.get("version")
+                        language_key_version = LanguageVersion(element.get("key"), element.get("version"))
                     else:
                         raise ValueError(
                             f"Language should be an object. Found: {element}"
@@ -292,11 +290,9 @@ def _deserialize_classifier_instance(
 
             for containment_entry in containments:
                 containment_obj = cast(JsonObject, containment_entry)
-                ids = SerializationUtils.try_to_get_array_of_ids(
+                ids : List[Optional[str]] = SerializationUtils.try_to_get_array_of_ids(
                     containment_obj, "children"
-                )
-                if ids is None:
-                    ids = []
+                ) or []
                 mp = SerializationUtils.try_to_get_meta_pointer_property(
                     containment_obj, "containment"
                 )

src/lionweb/serialization/protobuf_serialization.py

@@ -3,9 +3,8 @@
 import lionweb.serialization.proto.Chunk_pb2 as pb
 from lionweb.lionweb_version import LionWebVersion
 from lionweb.serialization import MetaPointer, SerializedClassifierInstance, SerializedPropertyValue, \
-    SerializedContainmentValue, SerializedReferenceValue
+    SerializedContainmentValue, SerializedReferenceValue, SerializationChunk
 from lionweb.serialization.data import LanguageVersion
-from lionweb.serialization.data import SerializationChunk
 from lionweb.serialization.data.serialized_reference_value import SerializedReferenceValueEntry
 from lionweb.serialization.deserialization_exception import DeserializationException
 
@@ -14,6 +13,7 @@
     from lionweb.serialization.data import SerializationChunk
 
 
+
 # from io.lionweb.serialization.data import (
 #     SerializationChunk,
 #     SerializedClassifierInstance,
@@ -68,12 +68,15 @@ def __init__(self, lionweb_version: Optional["LionWebVersion"] = LionWebVersion.
     #     return self._deserialize_serialization_chunk_to_instances(serialization_chunk)
     #
 
-    def read_chunk_from_bytes(self, data: bytes) -> pb.PBChunk:
+    def read_pbchunk_from_bytes(self, data: bytes) -> pb.PBChunk:
         """Read a protobuf Chunk from binary content"""
         self._chunk_instance.Clear()  # Reset the instance
         self._chunk_instance.ParseFromString(data)
         return self._chunk_instance
 
+    def read_chunk_from_bytes(self, data: bytes) -> SerializationChunk:
+        return self._deserialize_pbchunk_to_serialization_chunk(self.read_pbchunk_from_bytes(data))
+
 
     def _deserialize_pbchunk_to_serialization_chunk(self, chunk: pb.PBChunk) -> SerializationChunk:
         # Pre-size arrays as in Java
@@ -101,17 +104,19 @@ def _deserialize_pbchunk_to_serialization_chunk(self, chunk: pb.PBChunk) -> Seri
                     f"Unable to deserialize meta pointer with language {mp.li_language}"
                 )
             language_version = languages_array[mp.li_language]
+            language_key : Optional[str] = language_version.key if language_version is not None else None
+            language_v : Optional[str]= language_version.version if language_version is not None else None
             meta_pointer = MetaPointer(
-                language_version.key, language_version.version, strings_array[mp.si_key]
+                language_key, language_v, strings_array[mp.si_key]
             )
             metapointers_array[i] = meta_pointer
 
         serialization_chunk = SerializationChunk()
         serialization_chunk.serialization_format_version = chunk.serialization_format_version
 
-        for lv in languages_array:
-            if lv is not None:
-                serialization_chunk.add_language(lv)
+        valid_languages = [lv for lv in languages_array if lv is not None]
+        for lv in valid_languages:
+            serialization_chunk.add_language(lv)
 
         # Nodes
         for n in chunk.nodes:
@@ -131,7 +136,7 @@ def _deserialize_pbchunk_to_serialization_chunk(self, chunk: pb.PBChunk) -> Seri
 
             # containments
             for c in n.containments:
-                children: List[str] = []
+                children: List[Optional[str]] = []
                 for child_index in c.si_children:
                     if child_index == 0:
                         raise DeserializationException(
@@ -344,3 +349,4 @@ def _deserialize_pbchunk_to_serialization_chunk(self, chunk: pb.PBChunk) -> Seri
     #         chunk.interned_meta_pointers.append(pmp)
     #
     #     return chunk
+