From 82b8c7792e028a7ea5f855f2e7e1fe13bdabbfcb Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Thu, 12 Jan 2023 18:19:15 +0100 Subject: [PATCH] Allow to specify a limit to concurrent transmissions --- pyrogram/client.py | 17 +++++++++++++---- pyrogram/methods/advanced/save_file.py | 2 +- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/pyrogram/client.py b/pyrogram/client.py index 90aeb1f9..9de55621 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -172,6 +172,11 @@ class Client(Methods): Pass True to hide the password when typing it during the login. Defaults to False, because ``getpass`` (the library used) is known to be problematic in some terminal environments. + + max_concurrent_transmissions (``bool``, *optional*): + Set the maximum amount of concurrent transmissions (uploads & downloads). + A value that is too high may result in network related issues. + Defaults to 1. """ APP_VERSION = f"Pyrogram {__version__}" @@ -189,6 +194,8 @@ class Client(Methods): # Interval of seconds in which the updates watchdog will kick in UPDATES_WATCHDOG_INTERVAL = 5 * 60 + MAX_CONCURRENT_TRANSMISSIONS = 1 + mimetypes = MimeTypes() mimetypes.readfp(StringIO(mime_types)) @@ -217,7 +224,8 @@ class Client(Methods): no_updates: bool = None, takeout: bool = None, sleep_threshold: int = Session.SLEEP_THRESHOLD, - hide_password: bool = False + hide_password: bool = False, + max_concurrent_transmissions: int = MAX_CONCURRENT_TRANSMISSIONS ): super().__init__() @@ -245,6 +253,7 @@ class Client(Methods): self.takeout = takeout self.sleep_threshold = sleep_threshold self.hide_password = hide_password + self.max_concurrent_transmissions = max_concurrent_transmissions self.executor = ThreadPoolExecutor(self.workers, thread_name_prefix="Handler") @@ -266,8 +275,8 @@ class Client(Methods): self.media_sessions = {} self.media_sessions_lock = asyncio.Lock() - self.save_file_lock = asyncio.Lock() - self.get_file_lock = asyncio.Lock() + self.save_file_semaphore = asyncio.Semaphore(self.max_concurrent_transmissions) + self.get_file_semaphore = asyncio.Semaphore(self.max_concurrent_transmissions) self.is_connected = None self.is_initialized = None @@ -798,7 +807,7 @@ class Client(Methods): progress: Callable = None, progress_args: tuple = () ) -> Optional[AsyncGenerator[bytes, None]]: - async with self.get_file_lock: + async with self.get_file_semaphore: file_type = file_id.file_type if file_type == FileType.CHAT_PHOTO: diff --git a/pyrogram/methods/advanced/save_file.py b/pyrogram/methods/advanced/save_file.py index e683fe52..453a62af 100644 --- a/pyrogram/methods/advanced/save_file.py +++ b/pyrogram/methods/advanced/save_file.py @@ -94,7 +94,7 @@ class SaveFile: Raises: RPCError: In case of a Telegram RPC error. """ - async with self.save_file_lock: + async with self.save_file_semaphore: if path is None: return None