Rewrite save_file to handle missing file parts
This commit is contained in:
Dan
parent
b23b41bc7d
commit
bdc2ab6c73
@ -35,7 +35,7 @@ from pyrogram.api.errors import (
|
||||
PhoneMigrate, NetworkMigrate, PhoneNumberInvalid,
|
||||
PhoneNumberUnoccupied, PhoneCodeInvalid, PhoneCodeHashEmpty,
|
||||
PhoneCodeExpired, PhoneCodeEmpty, SessionPasswordNeeded,
|
||||
PasswordHashInvalid, FloodWait, PeerIdInvalid
|
||||
PasswordHashInvalid, FloodWait, PeerIdInvalid, FilePartMissing
|
||||
)
|
||||
from pyrogram.api.types import (
|
||||
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,
|
||||
chat_id: int or 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_total_parts = math.ceil(file_size / 512 / 1024)
|
||||
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)
|
||||
file_total_parts = math.ceil(file_size / part_size)
|
||||
# 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
|
||||
|
||||
try:
|
||||
session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
|
||||
session.start()
|
||||
|
||||
file_id = session.msg_id()
|
||||
md5_sum = md5()
|
||||
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
file_part = 0
|
||||
f.seek(part_size * file_part)
|
||||
|
||||
while True:
|
||||
chunk = f.read(512 * 1024)
|
||||
chunk = f.read(part_size)
|
||||
|
||||
if not chunk:
|
||||
if not is_big:
|
||||
md5_sum = "".join([hex(i)[2:].zfill(2) for i in md5_sum.digest()])
|
||||
break
|
||||
|
||||
md5_sum.update(chunk)
|
||||
|
||||
session.send(
|
||||
(functions.upload.SaveBigFilePart if is_big else functions.upload.SaveFilePart)(
|
||||
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
|
||||
except Exception as e:
|
||||
log.error(e)
|
||||
else:
|
||||
return types.InputFile(
|
||||
return (types.InputFileBig if is_big else types.InputFile)(
|
||||
id=file_id,
|
||||
parts=file_part,
|
||||
parts=file_total_parts,
|
||||
name=os.path.basename(path),
|
||||
md5_checksum=md5_sum
|
||||
)
|
||||
@ -566,11 +603,15 @@ class Client:
|
||||
ttl_seconds: int = None,
|
||||
disable_notification: bool = None,
|
||||
reply_to_message_id: int = None):
|
||||
return self.send(
|
||||
file = self.save_file(photo)
|
||||
|
||||
while True:
|
||||
try:
|
||||
r = self.send(
|
||||
functions.messages.SendMedia(
|
||||
peer=self.resolve_peer(chat_id),
|
||||
media=types.InputMediaUploadedPhoto(
|
||||
file=self.save_file(photo),
|
||||
file=file,
|
||||
caption=caption,
|
||||
ttl_seconds=ttl_seconds
|
||||
),
|
||||
@ -579,6 +620,10 @@ class Client:
|
||||
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,
|
||||
chat_id: int or str,
|
||||
@ -589,12 +634,16 @@ class Client:
|
||||
title: str = None,
|
||||
disable_notification: bool = None,
|
||||
reply_to_message_id: int = None):
|
||||
return self.send(
|
||||
file = self.save_file(audio)
|
||||
|
||||
while True:
|
||||
try:
|
||||
r = self.send(
|
||||
functions.messages.SendMedia(
|
||||
peer=self.resolve_peer(chat_id),
|
||||
media=types.InputMediaUploadedDocument(
|
||||
mime_type=mimetypes.types_map.get("." + audio.split(".")[-1], "audio/mpeg"),
|
||||
file=self.save_file(audio),
|
||||
file=file,
|
||||
caption=caption,
|
||||
attributes=[
|
||||
types.DocumentAttributeAudio(
|
||||
@ -610,6 +659,10 @@ class Client:
|
||||
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,
|
||||
chat_id: int or str,
|
||||
@ -617,12 +670,16 @@ class Client:
|
||||
caption: str = "",
|
||||
disable_notification: bool = None,
|
||||
reply_to_message_id: int = None):
|
||||
return self.send(
|
||||
file = self.save_file(document)
|
||||
|
||||
while True:
|
||||
try:
|
||||
r = self.send(
|
||||
functions.messages.SendMedia(
|
||||
peer=self.resolve_peer(chat_id),
|
||||
media=types.InputMediaUploadedDocument(
|
||||
mime_type=mimetypes.types_map.get("." + document.split(".")[-1], "text/plain"),
|
||||
file=self.save_file(document),
|
||||
file=file,
|
||||
caption=caption,
|
||||
attributes=[
|
||||
types.DocumentAttributeFilename(os.path.basename(document))
|
||||
@ -633,6 +690,10 @@ class Client:
|
||||
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,
|
||||
chat_id: int or str,
|
||||
@ -643,12 +704,16 @@ class Client:
|
||||
caption: str = "",
|
||||
disable_notification: bool = None,
|
||||
reply_to_message_id: int = None):
|
||||
return self.send(
|
||||
file = self.save_file(video)
|
||||
|
||||
while True:
|
||||
try:
|
||||
r = self.send(
|
||||
functions.messages.SendMedia(
|
||||
peer=self.resolve_peer(chat_id),
|
||||
media=types.InputMediaUploadedDocument(
|
||||
mime_type=mimetypes.types_map[".mp4"],
|
||||
file=self.save_file(video),
|
||||
file=file,
|
||||
caption=caption,
|
||||
attributes=[
|
||||
types.DocumentAttributeVideo(
|
||||
@ -663,6 +728,10 @@ class Client:
|
||||
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,
|
||||
chat_id: int or str,
|
||||
@ -671,12 +740,16 @@ class Client:
|
||||
duration: int = 0,
|
||||
disable_notification: bool = None,
|
||||
reply_to_message_id: int = None):
|
||||
return self.send(
|
||||
file = self.save_file(voice)
|
||||
|
||||
while True:
|
||||
try:
|
||||
r = self.send(
|
||||
functions.messages.SendMedia(
|
||||
peer=self.resolve_peer(chat_id),
|
||||
media=types.InputMediaUploadedDocument(
|
||||
mime_type=mimetypes.types_map.get("." + voice.split(".")[-1], "audio/mpeg"),
|
||||
file=self.save_file(voice),
|
||||
file=file,
|
||||
caption=caption,
|
||||
attributes=[
|
||||
types.DocumentAttributeAudio(
|
||||
@ -690,6 +763,10 @@ class Client:
|
||||
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,
|
||||
chat_id: int or str,
|
||||
@ -698,12 +775,16 @@ class Client:
|
||||
length: int = 1,
|
||||
disable_notification: bool = 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(
|
||||
peer=self.resolve_peer(chat_id),
|
||||
media=types.InputMediaUploadedDocument(
|
||||
mime_type=mimetypes.types_map[".mp4"],
|
||||
file=self.save_file(video_note),
|
||||
file=file,
|
||||
caption="",
|
||||
attributes=[
|
||||
types.DocumentAttributeVideo(
|
||||
@ -719,32 +800,7 @@ class Client:
|
||||
random_id=self.rnd_id()
|
||||
)
|
||||
)
|
||||
|
||||
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()
|
||||
)
|
||||
)
|
||||
except FilePartMissing as e:
|
||||
self.save_file(video_note, file_id=file.id, file_part=e.x)
|
||||
else:
|
||||
return r
|
||||
|
||||
Loading…
Reference in New Issue
Block a user