From c7985b5131ad68377ab246788d94d1b4a052da39 Mon Sep 17 00:00:00 2001
From: wulan17 <wulan17@nusantararom.org>
Date: Fri, 17 Nov 2023 21:52:58 +0700
Subject: [PATCH] Pyrofork: only load MongoStorage module when pymongo is
 available #32

Signed-off-by: wulan17 <wulan17@nusantararom.org>
---
 pyrogram/client.py           | 21 ++++++++++++++++++---
 pyrogram/storage/__init__.py |  7 ++++++-
 requirements.txt             |  1 -
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/pyrogram/client.py b/pyrogram/client.py
index f694b4f7..d516c0fe 100644
--- a/pyrogram/client.py
+++ b/pyrogram/client.py
@@ -50,7 +50,13 @@ from pyrogram.errors import (
 from pyrogram.handlers.handler import Handler
 from pyrogram.methods import Methods
 from pyrogram.session import Auth, Session
-from pyrogram.storage import FileStorage, MemoryStorage, MongoStorage, Storage
+from pyrogram.storage import FileStorage, MemoryStorage, Storage
+try:
+    import pymongo
+except Exception:
+    pass
+else:
+    from pyrogram.storage import MongoStorage
 from pyrogram.types import User, TermsOfService
 from pyrogram.utils import ainput
 from .dispatcher import Dispatcher
@@ -129,7 +135,7 @@ class Client(Methods):
         mongodb (``dict``, *optional*):
             Mongodb config as dict, e.g.: *dict(connection=async_pymongo.AsyncClient("mongodb://..."), remove_peers=False)*.
             Only applicable for new sessions.
-        
+
         storage (:obj:`~pyrogram.storage.Storage`, *optional*):
             Custom session storage.
 
@@ -281,7 +287,16 @@ class Client(Methods):
         elif self.in_memory:
             self.storage = MemoryStorage(self.name)
         elif self.mongodb:
-            self.storage = MongoStorage(self.name, **self.mongodb)
+            try:
+                import pymongo
+            except Exception:
+                log.warning(
+                    "pymongo is missing! "
+                    "Using MemoryStorage as session storage"
+                )
+                self.storage = MemoryStorage(self.name)
+            else:
+                self.storage = MongoStorage(self.name, **self.mongodb)
         else:
             self.storage = FileStorage(self.name, self.workdir)
 
diff --git a/pyrogram/storage/__init__.py b/pyrogram/storage/__init__.py
index 09ff0e86..ae5a3b33 100644
--- a/pyrogram/storage/__init__.py
+++ b/pyrogram/storage/__init__.py
@@ -18,5 +18,10 @@
 
 from .file_storage import FileStorage
 from .memory_storage import MemoryStorage
-from .mongo_storage import MongoStorage
+try:
+    import pymongo
+except Exception:
+    pass
+else:
+    from .mongo_storage import MongoStorage
 from .storage import Storage
diff --git a/requirements.txt b/requirements.txt
index 99f08c64..e832ea0b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,3 @@
 pyaes==1.6.1
-pymongo
 pymediainfo-pyrofork>=6.0.1,<7.0.0
 pysocks==1.7.1