Make proxy loading simpler

delivrance · delivrance · commit 6dcdeda2447b · 2018-04-13T12:30:13.000+02:00
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
@@ -54,15 +54,6 @@
 log = logging.getLogger(__name__)
 
 
-class Proxy:
-    def __init__(self, enabled: bool, hostname: str, port: int, username: str, password: str):
-        self.enabled = enabled
-        self.hostname = hostname
-        self.port = port
-        self.username = username
-        self.password = password
-
-
 class Client:
     """This class represents a Client, the main mean for interacting with Telegram.
     It exposes bot-like methods for an easy access to the API as well as a simple way to
@@ -134,7 +125,7 @@ def __init__(self,
                  session_name: str,
                  api_id: int or str = None,
                  api_hash: str = None,
-                 proxy: dict or Proxy = None,
+                 proxy: dict = None,
                  test_mode: bool = False,
                  phone_number: str = None,
                  phone_code: str or callable = None,
@@ -830,22 +821,17 @@ def load_config(self):
             else:
                 raise AttributeError("No API Key found")
 
-        if self.proxy is not None:
-            self.proxy = Proxy(
-                enabled=True,
-                hostname=self.proxy["hostname"],
-                port=int(self.proxy["port"]),
-                username=self.proxy.get("username", None),
-                password=self.proxy.get("password", None)
-            )
-        elif parser.has_section("proxy"):
-            self.proxy = Proxy(
-                enabled=parser.getboolean("proxy", "enabled"),
-                hostname=parser.get("proxy", "hostname"),
-                port=parser.getint("proxy", "port"),
-                username=parser.get("proxy", "username", fallback=None) or None,
-                password=parser.get("proxy", "password", fallback=None) or None
-            )
+        if self.proxy:
+            pass
+        else:
+            self.proxy = {}
+
+            if parser.has_section("proxy"):
+                self.proxy["enabled"] = parser.getboolean("proxy", "enabled")
+                self.proxy["hostname"] = parser.get("proxy", "hostname")
+                self.proxy["port"] = parser.getint("proxy", "port")
+                self.proxy["username"] = parser.get("proxy", "username", fallback=None) or None
+                self.proxy["password"] = parser.get("proxy", "password", fallback=None) or None
 
     def load_session(self, session_name):
         try:
diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py
@@ -32,7 +32,7 @@ class Connection:
         2: TCPIntermediate
     }
 
-    def __init__(self, address: tuple, proxy: type, mode: int = 1):
+    def __init__(self, address: tuple, proxy: dict, mode: int = 1):
         self.address = address
         self.proxy = proxy
         self.mode = self.MODES.get(mode, TCPAbridged)
diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
@@ -18,7 +18,6 @@
 
 import logging
 import socket
-from collections import namedtuple
 
 try:
     import socks
@@ -32,29 +31,25 @@
 
 log = logging.getLogger(__name__)
 
-Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
-
 
 class TCP(socks.socksocket):
-    def __init__(self, proxy: Proxy):
+    def __init__(self, proxy: dict):
         super().__init__()
         self.settimeout(10)
-        self.proxy_enabled = False
-
-        if proxy and proxy.enabled:
-            self.proxy_enabled = True
+        self.proxy_enabled = proxy.get("enabled", False)
 
+        if proxy and self.proxy_enabled:
             self.set_proxy(
                 proxy_type=socks.SOCKS5,
-                addr=proxy.hostname,
-                port=proxy.port,
-                username=proxy.username,
-                password=proxy.password
+                addr=proxy["hostname"],
+                port=proxy["port"],
+                username=proxy["username"],
+                password=proxy["password"]
             )
 
             log.info("Using proxy {}:{}".format(
-                proxy.hostname,
-                proxy.port
+                proxy["hostname"],
+                proxy["port"]
             ))
 
     def close(self):
diff --git a/pyrogram/connection/transport/tcp/tcp_abridged.py b/pyrogram/connection/transport/tcp/tcp_abridged.py
@@ -24,7 +24,7 @@
 
 
 class TCPAbridged(TCP):
-    def __init__(self, proxy: type):
+    def __init__(self, proxy: dict):
         super().__init__(proxy)
         self.is_first_packet = None
 
diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py
@@ -26,7 +26,7 @@
 
 
 class TCPFull(TCP):
-    def __init__(self, proxy: type):
+    def __init__(self, proxy: dict):
         super().__init__(proxy)
         self.seq_no = None
 
diff --git a/pyrogram/connection/transport/tcp/tcp_intermediate.py b/pyrogram/connection/transport/tcp/tcp_intermediate.py
@@ -25,7 +25,7 @@
 
 
 class TCPIntermediate(TCP):
-    def __init__(self, proxy: type):
+    def __init__(self, proxy: dict):
         super().__init__(proxy)
         self.is_first_packet = None
 
diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py
@@ -46,7 +46,7 @@ class Auth:
         16
     )
 
-    def __init__(self, dc_id: int, test_mode: bool, proxy: type):
+    def __init__(self, dc_id: int, test_mode: bool, proxy: dict):
         self.dc_id = dc_id
         self.test_mode = test_mode
 
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
@@ -86,7 +86,7 @@ class Session:
     def __init__(self,
                  dc_id: int,
                  test_mode: bool,
-                 proxy: type,
+                 proxy: dict,
                  auth_key: bytes,
                  api_id: int,
                  is_cdn: bool = False,

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ class Connection:`
`32`	`32`	`2: TCPIntermediate`
`33`	`33`	`}`
`34`	`34`
`35`		`- def __init__(self, address: tuple, proxy: type, mode: int = 1):`
	`35`	`+ def __init__(self, address: tuple, proxy: dict, mode: int = 1):`
`36`	`36`	`self.address = address`
`37`	`37`	`self.proxy = proxy`
`38`	`38`	`self.mode = self.MODES.get(mode, TCPAbridged)`
Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ class Auth:`
`46`	`46`	`16`
`47`	`47`	`)`
`48`	`48`
`49`		`- def __init__(self, dc_id: int, test_mode: bool, proxy: type):`
	`49`	`+ def __init__(self, dc_id: int, test_mode: bool, proxy: dict):`
`50`	`50`	`self.dc_id = dc_id`
`51`	`51`	`self.test_mode = test_mode`
`52`	`52`