Rewrite save_file to handle missing file parts

This commit is contained in:
Dan 2017-12-18 14:41:32 +01:00
parent b23b41bc7d
commit bdc2ab6c73

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
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
session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id) session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
session.start()
try: try:
session.start()
file_id = session.msg_id()
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:
md5_sum = "".join([hex(i)[2:].zfill(2) for i in md5_sum.digest()]) if not is_big:
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,19 +603,27 @@ 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)
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), while True:
media=types.InputMediaUploadedPhoto( try:
file=self.save_file(photo), r = self.send(
caption=caption, functions.messages.SendMedia(
ttl_seconds=ttl_seconds peer=self.resolve_peer(chat_id),
), media=types.InputMediaUploadedPhoto(
silent=disable_notification or None, file=file,
reply_to_msg_id=reply_to_message_id, caption=caption,
random_id=self.rnd_id() ttl_seconds=ttl_seconds
) ),
) silent=disable_notification or None,
reply_to_msg_id=reply_to_message_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,27 +634,35 @@ 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)
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), while True:
media=types.InputMediaUploadedDocument( try:
mime_type=mimetypes.types_map.get("." + audio.split(".")[-1], "audio/mpeg"), r = self.send(
file=self.save_file(audio), functions.messages.SendMedia(
caption=caption, peer=self.resolve_peer(chat_id),
attributes=[ media=types.InputMediaUploadedDocument(
types.DocumentAttributeAudio( mime_type=mimetypes.types_map.get("." + audio.split(".")[-1], "audio/mpeg"),
duration=duration, file=file,
performer=performer, caption=caption,
title=title attributes=[
types.DocumentAttributeAudio(
duration=duration,
performer=performer,
title=title
),
types.DocumentAttributeFilename(os.path.basename(audio))
]
), ),
types.DocumentAttributeFilename(os.path.basename(audio)) silent=disable_notification or None,
] reply_to_msg_id=reply_to_message_id,
), random_id=self.rnd_id()
silent=disable_notification or None, )
reply_to_msg_id=reply_to_message_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,22 +670,30 @@ 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)
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), while True:
media=types.InputMediaUploadedDocument( try:
mime_type=mimetypes.types_map.get("." + document.split(".")[-1], "text/plain"), r = self.send(
file=self.save_file(document), functions.messages.SendMedia(
caption=caption, peer=self.resolve_peer(chat_id),
attributes=[ media=types.InputMediaUploadedDocument(
types.DocumentAttributeFilename(os.path.basename(document)) mime_type=mimetypes.types_map.get("." + document.split(".")[-1], "text/plain"),
] file=file,
), caption=caption,
silent=disable_notification or None, attributes=[
reply_to_msg_id=reply_to_message_id, types.DocumentAttributeFilename(os.path.basename(document))
random_id=self.rnd_id() ]
) ),
) silent=disable_notification or None,
reply_to_msg_id=reply_to_message_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,26 +704,34 @@ 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)
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), while True:
media=types.InputMediaUploadedDocument( try:
mime_type=mimetypes.types_map[".mp4"], r = self.send(
file=self.save_file(video), functions.messages.SendMedia(
caption=caption, peer=self.resolve_peer(chat_id),
attributes=[ media=types.InputMediaUploadedDocument(
types.DocumentAttributeVideo( mime_type=mimetypes.types_map[".mp4"],
duration=duration, file=file,
w=width, caption=caption,
h=height attributes=[
) types.DocumentAttributeVideo(
] duration=duration,
), w=width,
silent=disable_notification or None, h=height
reply_to_msg_id=reply_to_message_id, )
random_id=self.rnd_id() ]
) ),
) silent=disable_notification or None,
reply_to_msg_id=reply_to_message_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,25 +740,33 @@ 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)
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id), while True:
media=types.InputMediaUploadedDocument( try:
mime_type=mimetypes.types_map.get("." + voice.split(".")[-1], "audio/mpeg"), r = self.send(
file=self.save_file(voice), functions.messages.SendMedia(
caption=caption, peer=self.resolve_peer(chat_id),
attributes=[ media=types.InputMediaUploadedDocument(
types.DocumentAttributeAudio( mime_type=mimetypes.types_map.get("." + voice.split(".")[-1], "audio/mpeg"),
voice=True, file=file,
duration=duration caption=caption,
) attributes=[
] types.DocumentAttributeAudio(
), voice=True,
silent=disable_notification or None, duration=duration
reply_to_msg_id=reply_to_message_id, )
random_id=self.rnd_id() ]
) ),
) silent=disable_notification or None,
reply_to_msg_id=reply_to_message_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,53 +775,32 @@ 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)
functions.messages.SendMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument(
mime_type=mimetypes.types_map[".mp4"],
file=self.save_file(video_note),
caption="",
attributes=[
types.DocumentAttributeVideo(
round_message=True,
duration=duration,
w=length,
h=length
)
]
),
silent=disable_notification or None,
reply_to_msg_id=reply_to_message_id,
random_id=self.rnd_id()
)
)
def send_venue(self, while True:
chat_id: int or str, try:
latitude: float, r = self.send(
longitude: float, functions.messages.SendMedia(
title: str, peer=self.resolve_peer(chat_id),
address: str, media=types.InputMediaUploadedDocument(
foursquare_id: str = "", mime_type=mimetypes.types_map[".mp4"],
disable_notification: bool = None, file=file,
reply_to_message_id: int = None): caption="",
return self.send( attributes=[
functions.messages.SendMedia( types.DocumentAttributeVideo(
peer=self.resolve_peer(chat_id), round_message=True,
media=types.InputMediaVenue( duration=duration,
geo_point=types.InputGeoPoint( w=length,
lat=latitude, h=length
long=longitude )
), ]
title=title, ),
address=address, silent=disable_notification or None,
provider="", reply_to_msg_id=reply_to_message_id,
venue_id=foursquare_id, random_id=self.rnd_id()
venue_type="" )
), )
silent=disable_notification or None, except FilePartMissing as e:
reply_to_msg_id=reply_to_message_id, self.save_file(video_note, file_id=file.id, file_part=e.x)
random_id=self.rnd_id() else:
) return r
)