compressedfhir 0.0.1__py3-none-any.whl

compressedfhir/utilities/compressed_dict/v1/compressed_dict.py

@@ -0,0 +1,635 @@
+import copy
+import json
+from collections.abc import KeysView, ValuesView, ItemsView, MutableMapping
+from contextlib import contextmanager
+from typing import Dict, Optional, Iterator, cast, List, Any, overload, OrderedDict
+
+import msgpack
+import zlib
+
+from compressedfhir.utilities.compressed_dict.v1.compressed_dict_access_error import (
+    CompressedDictAccessError,
+)
+from compressedfhir.utilities.compressed_dict.v1.compressed_dict_storage_mode import (
+    CompressedDictStorageMode,
+    CompressedDictStorageType,
+)
+from compressedfhir.utilities.fhir_json_encoder import FhirJSONEncoder
+
+
+class CompressedDict[K, V](MutableMapping[K, V]):
+    """
+    A dictionary-like class that supports flexible storage options.
+
+    It can store data in raw format, MessagePack format, or compressed MessagePack format.
+    """
+
+    # use slots to reduce memory usage
+    __slots__ = [
+        "_storage_mode",
+        "_working_dict",
+        "_raw_dict",
+        "_serialized_dict",
+        "_properties_to_cache",
+        "_cached_properties",
+        "_length",
+        "_transaction_depth",
+    ]
+
+    def __init__(
+        self,
+        *,
+        initial_dict: Dict[K, V] | OrderedDict[K, V] | None = None,
+        storage_mode: CompressedDictStorageMode,
+        properties_to_cache: List[K] | None,
+    ) -> None:
+        """
+        Initialize a dictionary with flexible storage options
+        Args:
+            initial_dict: Initial dictionary to store
+            storage_mode: Storage method for dictionary contents
+                - 'raw': Store as original dictionary
+                - 'msgpack': Store as MessagePack serialized bytes
+                - 'compressed_msgpack': Store as compressed MessagePack bytes
+        """
+        # Storage configuration
+        self._storage_mode: CompressedDictStorageMode = storage_mode
+
+        # Working copy of the dictionary during context
+        self._working_dict: Optional[OrderedDict[K, V]] = None
+
+        # Private storage options
+        self._raw_dict: OrderedDict[K, V] = OrderedDict[K, V]()
+        self._serialized_dict: Optional[bytes] = None
+
+        self._properties_to_cache: List[K] | None = properties_to_cache
+
+        self._cached_properties: Dict[K, V] = {}
+
+        self._length: int = 0
+
+        self._transaction_depth: int = 0
+
+        # Populate initial dictionary if provided
+        if initial_dict:
+            # Ensure we use an OrderedDict to maintain original order
+            initial_dict_ordered = (
+                initial_dict
+                if isinstance(initial_dict, OrderedDict)
+                else OrderedDict[K, V](initial_dict)
+            )
+            self.replace(value=initial_dict_ordered)
+
+    @contextmanager
+    def transaction(self) -> Iterator["CompressedDict[K, V]"]:
+        """
+        Context manager to safely access and modify the dictionary.
+
+        Deserializes the dictionary before entering the context
+        Serializes the dictionary after exiting the context
+
+        Raises:
+            CompressedDictAccessError: If methods are called outside the context
+        """
+        try:
+            self.start_transaction()
+
+            # Yield the working dictionary
+            yield self
+
+        finally:
+            self.end_transaction()
+
+    def start_transaction(self) -> "CompressedDict[K, V]":
+        """
+        Starts a transaction.  Use transaction() for a contextmanager for simpler usage.
+        """
+        # Increment transaction depth
+        self._transaction_depth += 1
+        # Ensure working dictionary is ready on first entry
+        if self._transaction_depth == 1:
+            self.ensure_working_dict()
+
+        return self
+
+    def end_transaction(self) -> "CompressedDict[K, V]":
+        """
+        Ends a transaction.  Use transaction() for a context_manager for simpler usage.
+
+        """
+        # Decrement transaction depth
+        self._transaction_depth -= 1
+        # Only update serialized dict when outermost transaction completes
+        if self._transaction_depth == 0:
+            self._update_serialized_dict(current_dict=self._working_dict)
+
+            # Clear the working dictionary
+            self._working_dict = None
+
+        return self
+
+    def ensure_working_dict(self) -> None:
+        """
+        Ensures that the working dictionary is initialized and deserialized.
+
+        """
+        if not self._working_dict:
+            self._working_dict = self.create_working_dict()
+
+    def create_working_dict(self) -> OrderedDict[K, V]:
+        working_dict: OrderedDict[K, V]
+        # Deserialize the dictionary before entering the context
+        if self._storage_mode.storage_type == "raw":
+            # For raw mode, create a deep copy of the existing dictionary
+            working_dict = self._raw_dict
+        else:
+            # For serialized modes, deserialize
+            working_dict = (
+                self._deserialize_dict(
+                    serialized_dict=self._serialized_dict,
+                    storage_type=self._storage_mode.storage_type,
+                )
+                if self._serialized_dict
+                else OrderedDict[K, V]()
+            )
+            assert isinstance(working_dict, OrderedDict)
+        return working_dict
+
+    @staticmethod
+    def _serialize_dict(
+        *, dictionary: OrderedDict[K, V], storage_type: CompressedDictStorageType
+    ) -> bytes:
+        """
+        Serialize entire dictionary using MessagePack
+
+        Args:
+            dictionary: Dictionary to serialize
+            storage_type: Storage type to use for serialization
+
+        Returns:
+            Serialized bytes
+        """
+        assert isinstance(dictionary, OrderedDict)
+        if storage_type == "compressed":
+            # Serialize to JSON and compress with zlib
+            json_str = json.dumps(
+                dictionary, separators=(",", ":"), cls=FhirJSONEncoder
+            )  # Most compact JSON representation
+            return zlib.compress(
+                json_str.encode("utf-8"), level=zlib.Z_BEST_COMPRESSION
+            )
+
+        # Serialize using MessagePack
+        packed = msgpack.packb(
+            dictionary,
+            use_bin_type=True,  # Preserve string/bytes distinction
+            use_single_float=True,  # More compact float representation
+        )
+
+        # Optional compression
+        if storage_type == "compressed_msgpack":
+            packed = zlib.compress(packed, level=zlib.Z_BEST_COMPRESSION)
+
+        return packed
+
+    @staticmethod
+    def _deserialize_dict(
+        *,
+        serialized_dict: bytes,
+        storage_type: CompressedDictStorageType,
+    ) -> OrderedDict[K, V]:
+        """
+        Deserialize entire dictionary from MessagePack
+
+        Args:
+            serialized_dict: Serialized dictionary bytes
+
+        Returns:
+            Deserialized dictionary
+        """
+        assert serialized_dict is not None, "Serialized dictionary cannot be None"
+
+        if storage_type == "compressed":
+            # Decompress and parse JSON
+            decompressed = zlib.decompress(serialized_dict)
+            decoded_text = decompressed.decode("utf-8")
+            # noinspection PyTypeChecker
+            decompressed_dict = json.loads(decoded_text, object_pairs_hook=OrderedDict)
+            assert isinstance(decompressed_dict, OrderedDict)
+            return cast(OrderedDict[K, V], decompressed_dict)
+
+        # Decompress if needed
+        to_unpack = (
+            zlib.decompress(serialized_dict)
+            if storage_type == "compressed_msgpack"
+            else serialized_dict
+        )
+
+        # Deserialize
+        unpacked_dict = msgpack.unpackb(
+            to_unpack,
+            raw=False,  # Convert to strings
+            strict_map_key=False,  # Handle potential key type variations
+        )
+        unpacked_dict = (
+            unpacked_dict
+            if isinstance(unpacked_dict, OrderedDict)
+            else OrderedDict[K, V](unpacked_dict)
+        )
+        assert isinstance(unpacked_dict, OrderedDict)
+        return cast(
+            OrderedDict[K, V],
+            unpacked_dict,
+        )
+
+    def _get_dict(self) -> OrderedDict[K, V]:
+        """
+        Get the dictionary, deserializing if necessary
+
+        Returns:
+            Current dictionary state
+        """
+
+        if self._working_dict is None:
+            raise CompressedDictAccessError(
+                "Dictionary access is only allowed within an transaction() block. "
+                "Use 'with compressed_dict.transaction() as d:' to access the dictionary."
+                f"You tried to access it with storage type {self._storage_mode.storage_type}."
+            )
+
+        if self._storage_mode.storage_type == "raw":
+            return self._raw_dict
+
+        # For non-raw modes, do not keep deserialized dict
+        return self._working_dict
+
+    def __getitem__(self, key: K) -> V:
+        """
+        Retrieve a value
+
+        Args:
+            key: Dictionary key
+
+        Returns:
+            Value associated with the key
+        """
+
+        if self._properties_to_cache and key in self._properties_to_cache:
+            return self._cached_properties[key]
+
+        if self._working_dict is None:
+            raise CompressedDictAccessError(
+                "Dictionary access is only allowed within an transaction() block. "
+                "Use 'with compressed_dict.transaction() as d:' to access the dictionary."
+            )
+        return self._get_dict()[key]
+
+    def __setitem__(self, key: K, value: V) -> None:
+        """
+        Set a value
+
+        Args:
+            key: Dictionary key
+            value: Value to store
+        """
+        if self._working_dict is None:
+            raise CompressedDictAccessError(
+                "Dictionary modification is only allowed within an transaction() block. "
+                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
+            )
+
+        # Update the working dictionary
+        self._working_dict[key] = value
+
+        if self._properties_to_cache and key in self._properties_to_cache:
+            # Update the cached properties if the key is in the list
+            self._cached_properties[key] = value
+
+    def _update_serialized_dict(self, current_dict: OrderedDict[K, V] | None) -> None:
+        if current_dict is None:
+            self._cached_properties.clear()
+            self._length = 0
+            self._serialized_dict = None
+            self._raw_dict = OrderedDict[K, V]()
+            return
+
+        if self._properties_to_cache:
+            for key in self._properties_to_cache:
+                if key in current_dict:
+                    self._cached_properties[key] = current_dict[key]
+
+        self._length = len(current_dict)
+
+        if self._working_dict is not None:
+            # If the working dictionary is None, initialize it
+            self._working_dict = current_dict
+
+        if self._transaction_depth == 0:
+            # If we're in a transaction,
+            # The serialized dictionary will be updated after the transaction
+            if self._storage_mode.storage_type == "raw":
+                self._raw_dict = current_dict
+            else:
+                self._serialized_dict = (
+                    self._serialize_dict(
+                        dictionary=current_dict,
+                        storage_type=self._storage_mode.storage_type,
+                    )
+                    if current_dict
+                    else None
+                )
+
+    def __delitem__(self, key: K) -> None:
+        """
+        Delete an item
+
+        Args:
+            key: Key to delete
+        """
+        if self._working_dict is None:
+            raise CompressedDictAccessError(
+                "Dictionary modification is only allowed within an transaction() block. "
+                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
+            )
+
+        del self._working_dict[key]
+
+    def __contains__(self, key: object) -> bool:
+        """
+        Check if a key exists
+
+        Args:
+            key: Key to check
+
+        Returns:
+            Whether the key exists
+        """
+        # first check if the key is in the cached properties
+        if self._properties_to_cache and key in self._properties_to_cache:
+            return self._cached_properties.__contains__(key)
+
+        return self._get_dict().__contains__(key)
+
+    def __len__(self) -> int:
+        """
+        Get the number of items
+
+        Returns:
+            Number of items in the dictionary
+        """
+        return self._length
+
+    def __iter__(self) -> Iterator[K]:
+        """
+        Iterate over keys
+
+        Returns:
+            Iterator of keys
+        """
+        if self._working_dict is None:
+            raise CompressedDictAccessError(
+                "Dictionary modification is only allowed within an transaction() block. "
+                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
+            )
+        return iter(self._get_dict())
+
+    def keys(self) -> KeysView[K]:
+        """
+        Get an iterator of keys
+
+        Returns:
+            Iterator of keys
+        """
+        return self._get_dict().keys()
+
+    def values(self) -> ValuesView[V]:
+        """
+        Get an iterator of values
+
+        Returns:
+            Iterator of values
+        """
+        return self._get_dict().values()
+
+    def items(self) -> ItemsView[K, V]:
+        """
+        Get an iterator of key-value pairs
+
+        Returns:
+            Iterator of key-value pairs
+        """
+        return self._get_dict().items()
+
+    def dict(self, *, remove_nulls: bool = True) -> OrderedDict[K, V]:
+        """
+        Convert to a standard dictionary
+
+        Returns:
+            Standard dictionary with all values
+        """
+        if self._working_dict:
+            return self._working_dict
+        else:
+            # if the working dict is None, return it but don't store it in the self._working_dict to keep memory low
+            return self.create_working_dict()
+
+    def __repr__(self) -> str:
+        """
+        String representation of the dictionary
+
+        Returns:
+            String representation
+        """
+        cached_property_list: List[str] = [
+            f"{k}={v}" for k, v in self._cached_properties.items()
+        ]
+        return (
+            (
+                f"CompressedDict(storage_type='{self._storage_mode.storage_type}', keys={self._length}"
+            )
+            + (", " if cached_property_list else "")
+            + ", ".join(cached_property_list)
+            + ")"
+        )
+
+    def replace(
+        self, *, value: Dict[K, V] | OrderedDict[K, V]
+    ) -> "CompressedDict[K, V]":
+        """
+        Replace the current dictionary with a new one
+
+        Args:
+            value: New dictionary to store
+
+        Returns:
+            Self
+        """
+        if not value:
+            self.clear()
+            return self
+
+        new_dict: OrderedDict[K, V] = (
+            value if isinstance(value, OrderedDict) else OrderedDict[K, V](value)
+        )
+        self._update_serialized_dict(current_dict=new_dict)
+        return self
+
+    def clear(self) -> None:
+        """
+        Clear the dictionary
+        """
+        self._update_serialized_dict(current_dict=None)
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Check equality with another dictionary
+
+        Args:
+            other: Dictionary to compare with (CompressedDict or plain dict)
+
+        Returns:
+            Whether the dictionaries are equal in keys and values
+        """
+        # If other is not a dictionary-like object, return False
+        if not isinstance(other, (CompressedDict, dict, OrderedDict)):
+            return False
+
+        # Get the dictionary representation of self
+        self_dict = self.dict()
+
+        # If other is a CompressedDict, use its _get_dict() method
+        if isinstance(other, CompressedDict):
+            other_dict = other.dict()
+        else:
+            # If other is a plain dict, use it directly
+            if isinstance(other, OrderedDict):
+                # If other is an OrderedDict, use it directly
+                other_dict = other
+            else:
+                other_dict = OrderedDict[K, V](other)
+
+        # Compare keys and values
+        # Check that all keys in both dictionaries match exactly
+        return set(self_dict.keys()) == set(other_dict.keys()) and all(
+            self_dict[key] == other_dict[key] for key in self_dict
+        )
+
+    @overload
+    def get(self, key: K) -> Optional[V]:
+        """
+        Get a value for an existing key
+
+        :param key: Key to retrieve
+        :return: Value or None if key is not found
+        """
+        ...
+
+    # noinspection PyMethodOverriding
+    @overload
+    def get[_T](self, key: K, /, default: V | _T) -> V | _T:
+        """
+        Get a value with a default
+
+        Args:
+            key: Key to retrieve
+            default: Default value if key is not found
+
+        Returns:
+            Value associated with the key or default
+        """
+        ...
+
+    def get[_T](self, key: K, default: V | _T | None = None) -> V | _T | None:
+        if key in self:
+            return self[key]
+        return default
+
+    def __deepcopy__(self, memo: Dict[int, Any]) -> "CompressedDict[K, V]":
+        """
+        Create a deep copy of the dictionary
+
+        Args:
+            memo: Memoization dictionary for deep copy
+
+        Returns:
+            Deep copy of the dictionary
+        """
+        # Create a new instance with the same storage mode
+        new_instance = CompressedDict(
+            initial_dict=copy.deepcopy(self.dict()),
+            storage_mode=self._storage_mode,
+            properties_to_cache=self._properties_to_cache,
+        )
+        return new_instance
+
+    def copy(self) -> "CompressedDict[K,V]":
+        """
+        Creates a copy of the BundleEntry object.
+
+        :return: A new BundleEntry object with the same attributes.
+        """
+        return copy.deepcopy(self)
+
+    def get_storage_mode(self) -> CompressedDictStorageMode:
+        """
+        Get the storage mode
+
+        Returns:
+            Storage mode
+        """
+        return self._storage_mode
+
+    @overload
+    def pop(self, key: K, /) -> V:
+        """
+        Remove and return a value for an existing key
+
+        :param key: Key to remove
+        :return: Removed value
+        :raises KeyError: If key is not found
+        """
+        ...
+
+    # noinspection PyMethodOverriding
+    @overload
+    def pop[_T](self, key: K, /, default: _T) -> V | _T:
+        """
+        Remove and return a value, or return default if key is not found
+
+        :param key: Key to remove
+        :param default: Default value to return if key is not found
+        :return: Removed value or default
+        """
+        ...
+
+    def pop[_T](self, key: K, /, default: V | _T | None = None) -> V | _T | None:
+        """
+        Remove and return a value
+
+        :param key: Key to remove
+        :param default: Optional default value if key is not found
+        :return: Removed value or default
+        :raises KeyError: If key is not found and no default is provided
+        """
+        if self._working_dict is None:
+            raise CompressedDictAccessError(
+                "Dictionary modification is only allowed within a transaction() block. "
+                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
+            )
+
+        # If no default is provided, use the standard dict.pop() behavior
+        if default is None:
+            return self._working_dict.pop(key)
+
+        return self._working_dict.pop(key, default)
+
+    def to_plain_dict(self) -> Dict[K, V]:
+        """
+        Get the plain dictionary representation
+
+        Returns:
+            Plain dictionary
+        """
+        return cast(
+            Dict[K, V], json.loads(json.dumps(self.dict(), cls=FhirJSONEncoder))
+        )

compressedfhir/utilities/compressed_dict/v1/compressed_dict_access_error.py

@@ -0,0 +1,2 @@
+class CompressedDictAccessError(Exception):
+    """Exception raised when attempting to access CompressedDict outside of context."""

compressedfhir/utilities/compressed_dict/v1/compressed_dict_storage_mode.py

@@ -0,0 +1,50 @@
+import dataclasses
+from typing import Literal, TypeAlias
+
+CompressedDictStorageType: TypeAlias = Literal[
+    "raw", "compressed", "msgpack", "compressed_msgpack"
+]
+"""
+CompressedDictStorageType is a type alias for the different storage types
+raw: No compression
+compressed: Compressed using zlib
+msgpack: Compressed using msgpack
+compressed_msgpack: Compressed using msgpack with zlib
+"""
+
+
+@dataclasses.dataclass(slots=True)
+class CompressedDictStorageMode:
+    """
+    CompressedDictStorageMode is a dataclass that defines the storage mode
+    """
+
+    storage_type: CompressedDictStorageType = "compressed"
+    """
+    CompressedDictStorageType is a type alias for the different storage types
+    raw: No compression
+    compressed: Compressed using zlib
+    msgpack: Compressed using msgpack
+    compressed_msgpack: Compressed using msgpack with zlib
+    """
+
+    @classmethod
+    def default(cls) -> "CompressedDictStorageMode":
+        """
+        Returns the default storage mode
+        """
+        return cls(storage_type="compressed")
+
+    @classmethod
+    def raw(cls) -> "CompressedDictStorageMode":
+        """
+        Returns the default storage mode
+        """
+        return cls(storage_type="raw")
+
+    @classmethod
+    def compressed(cls) -> "CompressedDictStorageMode":
+        """
+        Returns the default storage mode
+        """
+        return cls(storage_type="compressed")