Skip to content

Commit b3faf21

Browse files
delivrancedelivrance
committed
Rework in-memory uploads
1 parent de8f784 commit b3faf21

9 files changed

Lines changed: 260 additions & 126 deletions

File tree

‎pyrogram/client/client.py‎

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
# You should have received a copy of the GNU Lesser General Public License
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

19+
import io
1920
import logging
2021
import math
2122
import os
@@ -30,7 +31,7 @@
3031
from pathlib import Path
3132
from signal import signal, SIGINT, SIGTERM, SIGABRT
3233
from threading import Thread
33-
from typing import Union, List
34+
from typing import Union, List, BinaryIO
3435

3536
from pyrogram.api import functions, types
3637
from pyrogram.api.core import TLObject
@@ -39,9 +40,9 @@
3940
from pyrogram.client.methods.password.utils import compute_check
4041
from pyrogram.crypto import AES
4142
from pyrogram.errors import (
42-
PhoneMigrate, NetworkMigrate, SessionPasswordNeeded,
43-
FloodWait, PeerIdInvalid, VolumeLocNotFound, UserMigrate, ChannelPrivate, AuthBytesInvalid,
44-
BadRequest)
43+
PhoneMigrate, NetworkMigrate, SessionPasswordNeeded, PeerIdInvalid, VolumeLocNotFound, UserMigrate, ChannelPrivate,
44+
AuthBytesInvalid, BadRequest
45+
)
4546
from pyrogram.session import Auth, Session
4647
from .ext import utils, Syncer, BaseClient, Dispatcher
4748
from .methods import Methods
@@ -1713,7 +1714,7 @@ def resolve_peer(self, peer_id: Union[int, str]):
17131714

17141715
def save_file(
17151716
self,
1716-
path: str,
1717+
path: Union[str, BinaryIO],
17171718
file_id: int = None,
17181719
file_part: int = 0,
17191720
progress: callable = None,
@@ -1767,7 +1768,19 @@ def save_file(
17671768
RPCError: In case of a Telegram RPC error.
17681769
"""
17691770
part_size = 512 * 1024
1770-
file_size = os.path.getsize(path)
1771+
1772+
if isinstance(path, str):
1773+
fp = open(path, "rb")
1774+
elif isinstance(path, io.IOBase):
1775+
fp = path
1776+
else:
1777+
raise ValueError("Invalid file. Expected a file path as string or a binary (not text) file pointer")
1778+
1779+
file_name = fp.name
1780+
1781+
fp.seek(0, os.SEEK_END)
1782+
file_size = fp.tell()
1783+
fp.seek(0)
17711784

17721785
if file_size == 0:
17731786
raise ValueError("File size equals to 0 B")
@@ -1785,11 +1798,11 @@ def save_file(
17851798
session.start()
17861799

17871800
try:
1788-
with open(path, "rb") as f:
1789-
f.seek(part_size * file_part)
1801+
with fp:
1802+
fp.seek(part_size * file_part)
17901803

17911804
while True:
1792-
chunk = f.read(part_size)
1805+
chunk = fp.read(part_size)
17931806

17941807
if not chunk:
17951808
if not is_big:
@@ -1835,14 +1848,14 @@ def save_file(
18351848
return types.InputFileBig(
18361849
id=file_id,
18371850
parts=file_total_parts,
1838-
name=os.path.basename(path),
1851+
name=file_name,
18391852

18401853
)
18411854
else:
18421855
return types.InputFile(
18431856
id=file_id,
18441857
parts=file_total_parts,
1845-
name=os.path.basename(path),
1858+
name=file_name,
18461859
md5_checksum=md5_sum
18471860
)
18481861
finally:

‎pyrogram/client/methods/messages/send_animation.py‎

Lines changed: 36 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import os
2020
import re
21-
from typing import Union
21+
from typing import Union, BinaryIO
2222

2323
import pyrogram
2424
from pyrogram.api import functions, types
@@ -30,15 +30,15 @@ class SendAnimation(BaseClient):
3030
def send_animation(
3131
self,
3232
chat_id: Union[int, str],
33-
animation: str,
33+
animation: Union[str, BinaryIO],
3434
file_ref: str = None,
3535
caption: str = "",
3636
unsave: bool = False,
3737
parse_mode: Union[str, None] = object,
3838
duration: int = 0,
3939
width: int = 0,
4040
height: int = 0,
41-
thumb: str = None,
41+
thumb: Union[str, BinaryIO] = None,
4242
file_name: str = None,
4343
disable_notification: bool = None,
4444
reply_to_message_id: int = None,
@@ -60,11 +60,12 @@ def send_animation(
6060
For your personal cloud (Saved Messages) you can simply use "me" or "self".
6161
For a contact that exists in your Telegram address book you can use his phone number (str).
6262
63-
animation (``str``):
63+
animation (``str`` | ``BinaryIO``):
6464
Animation to send.
6565
Pass a file_id as string to send an animation that exists on the Telegram servers,
66-
pass an HTTP URL as a string for Telegram to get an animation from the Internet, or
67-
pass a file path as string to upload a new animation that exists on your local machine.
66+
pass an HTTP URL as a string for Telegram to get an animation from the Internet,
67+
pass a file path as string to upload a new animation that exists on your local machine, or
68+
pass a binary file-like object with its attribute ".name" set for in-memory uploads.
6869
6970
file_ref (``str``, *optional*):
7071
A valid file reference obtained by a recently fetched media message.
@@ -93,7 +94,7 @@ def send_animation(
9394
height (``int``, *optional*):
9495
Animation height.
9596
96-
thumb (``str``, *optional*):
97+
thumb (``str`` | ``BinaryIO``, *optional*):
9798
Thumbnail of the animation file sent.
9899
The thumbnail should be in JPEG format and less than 200 KB in size.
99100
A thumbnail's width and height should not exceed 320 pixels.
@@ -164,11 +165,36 @@ def progress(current, total):
164165
file = None
165166

166167
try:
167-
if os.path.isfile(animation):
168+
if isinstance(animation, str):
169+
if os.path.isfile(animation):
170+
thumb = None if thumb is None else self.save_file(thumb)
171+
file = self.save_file(animation, progress=progress, progress_args=progress_args)
172+
media = types.InputMediaUploadedDocument(
173+
mime_type=self.guess_mime_type(animation) or "video/mp4",
174+
file=file,
175+
thumb=thumb,
176+
attributes=[
177+
types.DocumentAttributeVideo(
178+
supports_streaming=True,
179+
duration=duration,
180+
w=width,
181+
h=height
182+
),
183+
types.DocumentAttributeFilename(file_name=file_name or os.path.basename(animation)),
184+
types.DocumentAttributeAnimated()
185+
]
186+
)
187+
elif re.match("^https?://", animation):
188+
media = types.InputMediaDocumentExternal(
189+
url=animation
190+
)
191+
else:
192+
media = utils.get_input_media_from_file_id(animation, file_ref, 10)
193+
else:
168194
thumb = None if thumb is None else self.save_file(thumb)
169195
file = self.save_file(animation, progress=progress, progress_args=progress_args)
170196
media = types.InputMediaUploadedDocument(
171-
mime_type=self.guess_mime_type(animation) or "video/mp4",
197+
mime_type=self.guess_mime_type(animation.name) or "video/mp4",
172198
file=file,
173199
thumb=thumb,
174200
attributes=[
@@ -178,16 +204,10 @@ def progress(current, total):
178204
w=width,
179205
h=height
180206
),
181-
types.DocumentAttributeFilename(file_name=file_name or os.path.basename(animation)),
207+
types.DocumentAttributeFilename(file_name=animation.name),
182208
types.DocumentAttributeAnimated()
183209
]
184210
)
185-
elif re.match("^https?://", animation):
186-
media = types.InputMediaDocumentExternal(
187-
url=animation
188-
)
189-
else:
190-
media = utils.get_input_media_from_file_id(animation, file_ref, 10)
191211

192212
while True:
193213
try:

‎pyrogram/client/methods/messages/send_audio.py‎

Lines changed: 34 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import os
2020
import re
21-
from typing import Union
21+
from typing import Union, BinaryIO
2222

2323
import pyrogram
2424
from pyrogram.api import functions, types
@@ -30,14 +30,14 @@ class SendAudio(BaseClient):
3030
def send_audio(
3131
self,
3232
chat_id: Union[int, str],
33-
audio: str,
33+
audio: Union[str, BinaryIO],
3434
file_ref: str = None,
3535
caption: str = "",
3636
parse_mode: Union[str, None] = object,
3737
duration: int = 0,
3838
performer: str = None,
3939
title: str = None,
40-
thumb: str = None,
40+
thumb: Union[str, BinaryIO] = None,
4141
file_name: str = None,
4242
disable_notification: bool = None,
4343
reply_to_message_id: int = None,
@@ -61,11 +61,12 @@ def send_audio(
6161
For your personal cloud (Saved Messages) you can simply use "me" or "self".
6262
For a contact that exists in your Telegram address book you can use his phone number (str).
6363
64-
audio (``str``):
64+
audio (``str`` | ``BinaryIO``):
6565
Audio file to send.
6666
Pass a file_id as string to send an audio file that exists on the Telegram servers,
67-
pass an HTTP URL as a string for Telegram to get an audio file from the Internet, or
68-
pass a file path as string to upload a new audio file that exists on your local machine.
67+
pass an HTTP URL as a string for Telegram to get an audio file from the Internet,
68+
pass a file path as string to upload a new audio file that exists on your local machine, or
69+
pass a binary file-like object with its attribute ".name" set for in-memory uploads.
6970
7071
file_ref (``str``, *optional*):
7172
A valid file reference obtained by a recently fetched media message.
@@ -90,7 +91,7 @@ def send_audio(
9091
title (``str``, *optional*):
9192
Track name.
9293
93-
thumb (``str``, *optional*):
94+
thumb (``str`` | ``BinaryIO``, *optional*):
9495
Thumbnail of the music file album cover.
9596
The thumbnail should be in JPEG format and less than 200 KB in size.
9697
A thumbnail's width and height should not exceed 320 pixels.
@@ -164,11 +165,34 @@ def progress(current, total):
164165
file = None
165166

166167
try:
167-
if os.path.isfile(audio):
168+
if isinstance(audio, str):
169+
if os.path.isfile(audio):
170+
thumb = None if thumb is None else self.save_file(thumb)
171+
file = self.save_file(audio, progress=progress, progress_args=progress_args)
172+
media = types.InputMediaUploadedDocument(
173+
mime_type=self.guess_mime_type(audio) or "audio/mpeg",
174+
file=file,
175+
thumb=thumb,
176+
attributes=[
177+
types.DocumentAttributeAudio(
178+
duration=duration,
179+
performer=performer,
180+
title=title
181+
),
182+
types.DocumentAttributeFilename(file_name=file_name or os.path.basename(audio))
183+
]
184+
)
185+
elif re.match("^https?://", audio):
186+
media = types.InputMediaDocumentExternal(
187+
url=audio
188+
)
189+
else:
190+
media = utils.get_input_media_from_file_id(audio, file_ref, 9)
191+
else:
168192
thumb = None if thumb is None else self.save_file(thumb)
169193
file = self.save_file(audio, progress=progress, progress_args=progress_args)
170194
media = types.InputMediaUploadedDocument(
171-
mime_type=self.guess_mime_type(audio) or "audio/mpeg",
195+
mime_type=self.guess_mime_type(audio.name) or "audio/mpeg",
172196
file=file,
173197
thumb=thumb,
174198
attributes=[
@@ -177,15 +201,9 @@ def progress(current, total):
177201
performer=performer,
178202
title=title
179203
),
180-
types.DocumentAttributeFilename(file_name=file_name or os.path.basename(audio))
204+
types.DocumentAttributeFilename(file_name=audio.name)
181205
]
182206
)
183-
elif re.match("^https?://", audio):
184-
media = types.InputMediaDocumentExternal(
185-
url=audio
186-
)
187-
else:
188-
media = utils.get_input_media_from_file_id(audio, file_ref, 9)
189207

190208
while True:
191209
try:

‎pyrogram/client/methods/messages/send_document.py‎

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import os
2020
import re
21-
from typing import Union
21+
from typing import Union, BinaryIO
2222

2323
import pyrogram
2424
from pyrogram.api import functions, types
@@ -30,9 +30,9 @@ class SendDocument(BaseClient):
3030
def send_document(
3131
self,
3232
chat_id: Union[int, str],
33-
document: str,
33+
document: Union[str, BinaryIO],
3434
file_ref: str = None,
35-
thumb: str = None,
35+
thumb: Union[str, BinaryIO] = None,
3636
caption: str = "",
3737
parse_mode: Union[str, None] = object,
3838
file_name: str = None,
@@ -56,17 +56,18 @@ def send_document(
5656
For your personal cloud (Saved Messages) you can simply use "me" or "self".
5757
For a contact that exists in your Telegram address book you can use his phone number (str).
5858
59-
document (``str``):
59+
document (``str`` | ``BinaryIO``):
6060
File to send.
6161
Pass a file_id as string to send a file that exists on the Telegram servers,
62-
pass an HTTP URL as a string for Telegram to get a file from the Internet, or
63-
pass a file path as string to upload a new file that exists on your local machine.
62+
pass an HTTP URL as a string for Telegram to get a file from the Internet,
63+
pass a file path as string to upload a new file that exists on your local machine, or
64+
pass a binary file-like object with its attribute ".name" set for in-memory uploads.
6465
6566
file_ref (``str``, *optional*):
6667
A valid file reference obtained by a recently fetched media message.
6768
To be used in combination with a file id in case a file reference is needed.
6869
69-
thumb (``str``, *optional*):
70+
thumb (``str`` | ``BinaryIO``, *optional*):
7071
Thumbnail of the file sent.
7172
The thumbnail should be in JPEG format and less than 200 KB in size.
7273
A thumbnail's width and height should not exceed 320 pixels.
@@ -144,23 +145,35 @@ def progress(current, total):
144145
file = None
145146

146147
try:
147-
if os.path.isfile(document):
148+
if isinstance(document, str):
149+
if os.path.isfile(document):
150+
thumb = None if thumb is None else self.save_file(thumb)
151+
file = self.save_file(document, progress=progress, progress_args=progress_args)
152+
media = types.InputMediaUploadedDocument(
153+
mime_type=self.guess_mime_type(document) or "application/zip",
154+
file=file,
155+
thumb=thumb,
156+
attributes=[
157+
types.DocumentAttributeFilename(file_name=file_name or os.path.basename(document))
158+
]
159+
)
160+
elif re.match("^https?://", document):
161+
media = types.InputMediaDocumentExternal(
162+
url=document
163+
)
164+
else:
165+
media = utils.get_input_media_from_file_id(document, file_ref, 5)
166+
else:
148167
thumb = None if thumb is None else self.save_file(thumb)
149168
file = self.save_file(document, progress=progress, progress_args=progress_args)
150169
media = types.InputMediaUploadedDocument(
151-
mime_type=self.guess_mime_type(document) or "application/zip",
170+
mime_type=self.guess_mime_type(document.name) or "application/zip",
152171
file=file,
153172
thumb=thumb,
154173
attributes=[
155-
types.DocumentAttributeFilename(file_name=file_name or os.path.basename(document))
174+
types.DocumentAttributeFilename(file_name=document.name)
156175
]
157176
)
158-
elif re.match("^https?://", document):
159-
media = types.InputMediaDocumentExternal(
160-
url=document
161-
)
162-
else:
163-
media = utils.get_input_media_from_file_id(document, file_ref, 5)
164177

165178
while True:
166179
try:

0 commit comments

Comments
 (0)