From c69acf78994e19670ce7f2fa6a319c91e314660f Mon Sep 17 00:00:00 2001
From: bhimrazy <bhimrajyadav977@gmail.com>
Date: Fri, 12 Jul 2024 02:32:47 +0545
Subject: [PATCH] chore: adds cryptography as conditional import

---
 src/litdata/constants.py            |  1 +
 src/litdata/utilities/encryption.py | 11 +++++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/litdata/constants.py b/src/litdata/constants.py
index 8befa208..e09e0b65 100644
--- 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
index 90cace46..f285cb43 100644
--- 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):