Don't use sys.exit(), re-raise ImportError instead

delivrance · delivrance · commit 42a287884258 · 2018-03-26T03:39:30.000+02:00
diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
@@ -18,17 +18,18 @@
 
 import logging
 import socket
-import sys
 from collections import namedtuple
 
 try:
     import socks
-except ImportError:
-    sys.exit(
+except ImportError as e:
+    e.msg = (
         "PySocks is missing and Pyrogram can't run without. "
         "Please install it using \"pip3 install pysocks\"."
     )
 
+    raise e
+
 log = logging.getLogger(__name__)
 
 Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
diff --git a/pyrogram/crypto/aes.py b/pyrogram/crypto/aes.py
@@ -16,19 +16,18 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-import sys
-
 try:
     import tgcrypto
-except ImportError:
-    sys.exit(
+except ImportError as e:
+    e.msg = (
         "TgCrypto is missing and Pyrogram can't run without. "
         "Please install it using \"pip3 install tgcrypto\". "
         "More info: https://docs.pyrogram.ml/resources/TgCrypto"
     )
 
+    raise e
+
 
-# TODO: Ugly IFs
 class AES:
     @classmethod
     def ige_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
