chore: adds cryptography as conditional import

Lightning-AI · Jul 11, 2024 · c69acf7 · c69acf7
1 parent a2d80c9
commit c69acf7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/litdata/constants.py b/src/litdata/constants.py
@@ -31,6 +31,7 @@
 _BOTO3_AVAILABLE = RequirementCache("boto3")
 _TORCH_AUDIO_AVAILABLE = RequirementCache("torchaudio")
 _ZSTD_AVAILABLE = RequirementCache("zstd")
+_CRYPTOGRAPHY_AVAILABLE = RequirementCache("cryptography")
 _GOOGLE_STORAGE_AVAILABLE = RequirementCache("google.cloud.storage")
 _TQDM_AVAILABLE = RequirementCache("tqdm")
 

diff --git a/src/litdata/utilities/encryption.py b/src/litdata/utilities/encryption.py
@@ -3,10 +3,13 @@
 from abc import ABC, abstractmethod
 from typing import Optional
 
-from cryptography.fernet import Fernet
-from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import padding, rsa
-from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from litdata.constants import _CRYPTOGRAPHY_AVAILABLE
+
+if _CRYPTOGRAPHY_AVAILABLE:
+    from cryptography.fernet import Fernet
+    from cryptography.hazmat.primitives import hashes, serialization
+    from cryptography.hazmat.primitives.asymmetric import padding, rsa
+    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 
 class Encryption(ABC):