fix: Updates encryption level handlings with enums

src/litdata/streaming/item_loader.py

@@ -169,7 +169,7 @@ def _load_encrypted_data(
         self._validate_encryption(encryption)
 
         # chunk-level decryption
-        if self._config["encryption"]["level"] == EncryptionLevel.CHUNK.value:
+        if self._config["encryption"]["level"] == EncryptionLevel.CHUNK:
             decrypted_data = self._decrypted_chunks.get(chunk_index, None)
             if decrypted_data is None:
                 with open(chunk_filepath, "rb", 0) as fp:
@@ -181,7 +181,7 @@ def _load_encrypted_data(
             data = self._load_data(BytesIO(decrypted_data), offset)
 
         # sample-level decryption
-        elif self._config["encryption"]["level"] == EncryptionLevel.SAMPLE.value:
+        elif self._config["encryption"]["level"] == EncryptionLevel.SAMPLE:
             with open(chunk_filepath, "rb", 0) as fp:
                 data = self._load_data(fp, offset)
                 data = encryption.decrypt(data)  # type: ignore

src/litdata/streaming/writer.py

@@ -27,7 +27,7 @@
 from litdata.streaming.compression import _COMPRESSORS, Compressor
 from litdata.streaming.serializers import Serializer, _get_serializers
 from litdata.utilities._pytree import PyTree, tree_flatten, treespec_dumps
-from litdata.utilities.encryption import Encryption
+from litdata.utilities.encryption import Encryption, EncryptionLevel
 from litdata.utilities.env import _DistributedEnv, _WorkerEnv
 from litdata.utilities.format import _convert_bytes_to_int, _human_readable_bytes
 
@@ -244,7 +244,7 @@ def _create_chunk(self, filename: str, on_done: bool = False) -> bytes:
         current_chunk_bytes = sum([item.bytes for item in items])
 
         # Whether to encrypt the data at the chunk level
-        if self._encryption and self._encryption.level == "chunk":
+        if self._encryption and self._encryption.level == EncryptionLevel.CHUNK:
             data = self._encryption.encrypt(data)
             current_chunk_bytes = len(data)
 
@@ -305,7 +305,7 @@ def add_item(self, index: int, items: Any) -> Optional[str]:
         data, dim = self.serialize(items)
 
         # Whether to encrypt the data at the sample level
-        if self._encryption and self._encryption.level == "sample":
+        if self._encryption and self._encryption.level == EncryptionLevel.SAMPLE:
             data = self._encryption.encrypt(data)
 
         self._serialized_items[index] = Item(

src/litdata/utilities/encryption.py

@@ -2,7 +2,7 @@
 import json
 import os
 from abc import ABC, abstractmethod
-from enum import Enum
+from dataclasses import dataclass
 from typing import Any, Dict, Literal, Tuple, Union, get_args
 
 from litdata.constants import _CRYPTOGRAPHY_AVAILABLE
@@ -14,7 +14,8 @@
     from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 
-class EncryptionLevel(Enum):
+@dataclass
+class EncryptionLevel:
     SAMPLE = "sample"
     CHUNK = "chunk"
 
@@ -62,7 +63,7 @@ class FernetEncryption(Encryption):
     def __init__(
         self,
         password: str,
-        level: EncryptionLevelType = EncryptionLevel.SAMPLE.value,
+        level: EncryptionLevelType = "sample",
     ) -> None:
         super().__init__()
         if not _CRYPTOGRAPHY_AVAILABLE:
@@ -131,7 +132,7 @@ class RSAEncryption(Encryption):
     def __init__(
         self,
         password: str,
-        level: EncryptionLevelType = EncryptionLevel.SAMPLE.value,
+        level: EncryptionLevelType = "sample",
     ) -> None:
         if not _CRYPTOGRAPHY_AVAILABLE:
             raise ModuleNotFoundError(str(_CRYPTOGRAPHY_AVAILABLE))