TeamPGM/pyrogram
2
0
Fork 0
mirror of https://github.com/TeamPGM/pyrogram.git synced 2024-11-24 07:51:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rewrite save_file to handle missing file parts

Browse Source
This commit is contained in:
Dan 2017-12-18 14:41:32 +01:00
parent b23b41bc7d
commit bdc2ab6c73
1 changed files with 209 additions and 153 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

168
pyrogram/client/client.py
View File

@ -35,7 +35,7 @@ from pyrogram.api.errors import (
PhoneMigrate, NetworkMigrate, PhoneNumberInvalid, PhoneMigrate, NetworkMigrate, PhoneNumberInvalid,
PhoneNumberUnoccupied, PhoneCodeInvalid, PhoneCodeHashEmpty, PhoneNumberUnoccupied, PhoneCodeInvalid, PhoneCodeHashEmpty,
PhoneCodeExpired, PhoneCodeEmpty, SessionPasswordNeeded, PhoneCodeExpired, PhoneCodeEmpty, SessionPasswordNeeded,
PasswordHashInvalid, FloodWait, PeerIdInvalid PasswordHashInvalid, FloodWait, PeerIdInvalid, FilePartMissing
) )
from pyrogram.api.types import ( from pyrogram.api.types import (
User, Chat, Channel, User, Chat, Channel,
@ -454,6 +454,35 @@ class Client:
) )
) )
def send_venue(self,
chat_id: int or str,
latitude: float,
longitude: float,
title: str,
address: str,
foursquare_id: str = "",
disable_notification: bool = None,
reply_to_message_id: int = None):
return self.send(
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaVenue(
geo_point=types.InputGeoPoint(
lat=latitude,
long=longitude
),
title=title,
address=address,
provider="",
venue_id=foursquare_id,
venue_type=""
),
silent=disable_notification or None,
reply_to_msg_id=reply_to_message_id,
random_id=self.rnd_id()
)
)
def send_contact(self, def send_contact(self,
chat_id: int or str, chat_id: int or str,
phone_number: str, phone_number: str,
@ -513,30 +542,32 @@ class Client:
) )
) )
def save_file(self, path): # TODO: Remove redundant code
def save_file(self, path: str, file_id: int = None, file_part: int = 0):
part_size = 512 * 1024
file_size = os.path.getsize(path) file_size = os.path.getsize(path)
file_total_parts = math.ceil(file_size / 512 / 1024) file_total_parts = math.ceil(file_size / part_size)
is_big = True if file_size >= 10 * 1024 * 1024 else False # is_big = True if file_size > 10 * 1024 * 1024 else False
session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id) is_big = False # Treat all files as not-big to have the server check for the md5 sum
is_missing_part = True if file_id is not None else False
file_id = file_id or self.rnd_id()
md5_sum = md5() if not is_big and not is_missing_part else None
try: session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
session.start() session.start()
file_id = session.msg_id() try:
md5_sum = md5()
with open(path, "rb") as f: with open(path, "rb") as f:
file_part = 0 f.seek(part_size * file_part)
while True: while True:
chunk = f.read(512 * 1024) chunk = f.read(part_size)
if not chunk: if not chunk:
if not is_big:
md5_sum = "".join([hex(i)[2:].zfill(2) for i in md5_sum.digest()]) md5_sum = "".join([hex(i)[2:].zfill(2) for i in md5_sum.digest()])
break break
md5_sum.update(chunk)
session.send( session.send(
(functions.upload.SaveBigFilePart if is_big else functions.upload.SaveFilePart)( (functions.upload.SaveBigFilePart if is_big else functions.upload.SaveFilePart)(
file_id=file_id, file_id=file_id,
@ -546,13 +577,19 @@ class Client:
) )
) )
if is_missing_part:
return
if not is_big:
md5_sum.update(chunk)
file_part += 1 file_part += 1
except Exception as e: except Exception as e:
log.error(e) log.error(e)
else: else:
return types.InputFile( return (types.InputFileBig if is_big else types.InputFile)(
id=file_id, id=file_id,
parts=file_part, parts=file_total_parts,
name=os.path.basename(path), name=os.path.basename(path),
md5_checksum=md5_sum md5_checksum=md5_sum
) )
@ -566,11 +603,15 @@ class Client:
ttl_seconds: int = None, ttl_seconds: int = None,
disable_notification: bool = None, disable_notification: bool = None,
reply_to_message_id: int = None): reply_to_message_id: int = None):
return self.send( file = self.save_file(photo)
while True:
try:
r = self.send(
functions.messages.SendMedia( functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedPhoto( media=types.InputMediaUploadedPhoto(
file=self.save_file(photo), file=file,
caption=caption, caption=caption,
ttl_seconds=ttl_seconds ttl_seconds=ttl_seconds
), ),
@ -579,6 +620,10 @@ class Client:
random_id=self.rnd_id() random_id=self.rnd_id()
) )
) )
except FilePartMissing as e:
self.save_file(photo, file_id=file.id, file_part=e.x)
else:
return r
def send_audio(self, def send_audio(self,
chat_id: int or str, chat_id: int or str,
@ -589,12 +634,16 @@ class Client:
title: str = None, title: str = None,
disable_notification: bool = None, disable_notification: bool = None,
reply_to_message_id: int = None): reply_to_message_id: int = None):
return self.send( file = self.save_file(audio)
while True:
try:
r = self.send(
functions.messages.SendMedia( functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument( media=types.InputMediaUploadedDocument(
mime_type=mimetypes.types_map.get("." + audio.split(".")[-1], "audio/mpeg"), mime_type=mimetypes.types_map.get("." + audio.split(".")[-1], "audio/mpeg"),
file=self.save_file(audio), file=file,
caption=caption, caption=caption,
attributes=[ attributes=[
types.DocumentAttributeAudio( types.DocumentAttributeAudio(
@ -610,6 +659,10 @@ class Client:
random_id=self.rnd_id() random_id=self.rnd_id()
) )
) )
except FilePartMissing as e:
self.save_file(audio, file_id=file.id, file_part=e.x)
else:
return r
def send_document(self, def send_document(self,
chat_id: int or str, chat_id: int or str,
@ -617,12 +670,16 @@ class Client:
caption: str = "", caption: str = "",
disable_notification: bool = None, disable_notification: bool = None,
reply_to_message_id: int = None): reply_to_message_id: int = None):
return self.send( file = self.save_file(document)
while True:
try:
r = self.send(
functions.messages.SendMedia( functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument( media=types.InputMediaUploadedDocument(
mime_type=mimetypes.types_map.get("." + document.split(".")[-1], "text/plain"), mime_type=mimetypes.types_map.get("." + document.split(".")[-1], "text/plain"),
file=self.save_file(document), file=file,
caption=caption, caption=caption,
attributes=[ attributes=[
types.DocumentAttributeFilename(os.path.basename(document)) types.DocumentAttributeFilename(os.path.basename(document))
@ -633,6 +690,10 @@ class Client:
random_id=self.rnd_id() random_id=self.rnd_id()
) )
) )
except FilePartMissing as e:
self.save_file(document, file_id=file.id, file_part=e.x)
else:
return r
def send_video(self, def send_video(self,
chat_id: int or str, chat_id: int or str,
@ -643,12 +704,16 @@ class Client:
caption: str = "", caption: str = "",
disable_notification: bool = None, disable_notification: bool = None,
reply_to_message_id: int = None): reply_to_message_id: int = None):
return self.send( file = self.save_file(video)
while True:
try:
r = self.send(
functions.messages.SendMedia( functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument( media=types.InputMediaUploadedDocument(
mime_type=mimetypes.types_map[".mp4"], mime_type=mimetypes.types_map[".mp4"],
file=self.save_file(video), file=file,
caption=caption, caption=caption,
attributes=[ attributes=[
types.DocumentAttributeVideo( types.DocumentAttributeVideo(
@ -663,6 +728,10 @@ class Client:
random_id=self.rnd_id() random_id=self.rnd_id()
) )
) )
except FilePartMissing as e:
self.save_file(video, file_id=file.id, file_part=e.x)
else:
return r
def send_voice(self, def send_voice(self,
chat_id: int or str, chat_id: int or str,
@ -671,12 +740,16 @@ class Client:
duration: int = 0, duration: int = 0,
disable_notification: bool = None, disable_notification: bool = None,
reply_to_message_id: int = None): reply_to_message_id: int = None):
return self.send( file = self.save_file(voice)
while True:
try:
r = self.send(
functions.messages.SendMedia( functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument( media=types.InputMediaUploadedDocument(
mime_type=mimetypes.types_map.get("." + voice.split(".")[-1], "audio/mpeg"), mime_type=mimetypes.types_map.get("." + voice.split(".")[-1], "audio/mpeg"),
file=self.save_file(voice), file=file,
caption=caption, caption=caption,
attributes=[ attributes=[
types.DocumentAttributeAudio( types.DocumentAttributeAudio(
@ -690,6 +763,10 @@ class Client:
random_id=self.rnd_id() random_id=self.rnd_id()
) )
) )
except FilePartMissing as e:
self.save_file(voice, file_id=file.id, file_part=e.x)
else:
return r
def send_video_note(self, def send_video_note(self,
chat_id: int or str, chat_id: int or str,
@ -698,12 +775,16 @@ class Client:
length: int = 1, length: int = 1,
disable_notification: bool = None, disable_notification: bool = None,
reply_to_message_id: int = None): reply_to_message_id: int = None):
return self.send( file = self.save_file(video_note)
while True:
try:
r = self.send(
functions.messages.SendMedia( functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument( media=types.InputMediaUploadedDocument(
mime_type=mimetypes.types_map[".mp4"], mime_type=mimetypes.types_map[".mp4"],
file=self.save_file(video_note), file=file,
caption="", caption="",
attributes=[ attributes=[
types.DocumentAttributeVideo( types.DocumentAttributeVideo(
@ -719,32 +800,7 @@ class Client:
random_id=self.rnd_id() random_id=self.rnd_id()
) )
) )
except FilePartMissing as e:
def send_venue(self, self.save_file(video_note, file_id=file.id, file_part=e.x)
chat_id: int or str, else:
latitude: float, return r
longitude: float,
title: str,
address: str,
foursquare_id: str = "",
disable_notification: bool = None,
reply_to_message_id: int = None):
return self.send(
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaVenue(
geo_point=types.InputGeoPoint(
lat=latitude,
long=longitude
),
title=title,
address=address,
provider="",
venue_id=foursquare_id,
venue_type=""
),
silent=disable_notification or None,
reply_to_msg_id=reply_to_message_id,
random_id=self.rnd_id()
)
)