telegram_upload/download_files.py
import os
import sys
from typing import Iterable, Iterator, Optional, BinaryIO
from telethon.tl.types import Message, DocumentAttributeFilename
if sys.version_info < (3, 8):
cached_property = property
else:
from functools import cached_property
CHUNK_FILE_SIZE = 1024 * 1024
def pipe_file(read_file_name: str, write_file: BinaryIO):
"""Read a file by its file name and write in another file already open."""
with open(read_file_name, "rb") as read_file:
while True:
data = read_file.read(CHUNK_FILE_SIZE)
if data:
write_file.write(data)
else:
break
class JoinStrategyBase:
"""Base class to inherit join strategies. The strategies depend on the file type.
For example, zip files and rar files do not merge in the same way.
"""
def __init__(self):
self.download_files = []
def is_part(self, download_file: 'DownloadFile') -> bool:
"""Returns if the download file is part of this bundle."""
raise NotImplementedError
def add_download_file(self, download_file: 'DownloadFile') -> None:
"""Add a download file to this bundle."""
if download_file in self.download_files:
return
self.download_files.append(download_file)
@classmethod
def is_applicable(cls, download_file: 'DownloadFile') -> bool:
"""Returns if this strategy is applicable to the download file."""
raise NotImplementedError
def join_download_files(self):
"""Join the downloaded files in the bundle."""
raise NotImplementedError
class UnionJoinStrategy(JoinStrategyBase):
"""Join separate files without any application. These files have extension
01, 02, 03...
"""
base_name: Optional[str] = None
@staticmethod
def get_base_name(download_file: 'DownloadFile'):
"""Returns the file name without extension."""
return download_file.file_name.rsplit(".", 1)[0]
def add_download_file(self, download_file: 'DownloadFile') -> None:
"""Add a download file to this bundle."""
if self.base_name is None:
self.base_name = self.get_base_name(download_file)
super().add_download_file(download_file)
def is_part(self, download_file: 'DownloadFile') -> bool:
"""Returns if the download file is part of this bundle."""
return self.base_name == self.get_base_name(download_file)
@classmethod
def is_applicable(cls, download_file: 'DownloadFile') -> bool:
"""Returns if this strategy is applicable to the download file."""
return download_file.file_name_extension.isdigit()
def join_download_files(self):
"""Join the downloaded files in the bundle."""
download_files = self.download_files
sorted_files = sorted(download_files, key=lambda x: x.file_name_extension)
sorted_files = [file for file in sorted_files if os.path.lexists(file.downloaded_file_name or "")]
if not sorted_files or len(sorted_files) - 1 != int(sorted_files[-1].file_name_extension):
# There are parts of the file missing. Stopping...
return
with open(self.get_base_name(sorted_files[0]), "wb") as new_file:
for download_file in sorted_files:
pipe_file(download_file.downloaded_file_name, new_file)
for download_file in sorted_files:
os.remove(download_file.downloaded_file_name)
JOIN_STRATEGIES = [
UnionJoinStrategy,
]
def get_join_strategy(download_file: 'DownloadFile') -> Optional[JoinStrategyBase]:
"""Get join strategy for the download file. An instance is returned if a strategy
is available. Otherwise, None is returned.
"""
for strategy_cls in JOIN_STRATEGIES:
if strategy_cls.is_applicable(download_file):
strategy = strategy_cls()
strategy.add_download_file(download_file)
return strategy
class DownloadFile:
"""File to download. This includes the Telethon message with the file."""
downloaded_file_name: Optional[str] = None
def __init__(self, message: Message):
"""Creates the download file instance from the message."""
self.message = message
def set_download_file_name(self, file_name):
"""After download the file, set the final download file name."""
self.downloaded_file_name = file_name
@cached_property
def filename_attr(self) -> Optional[DocumentAttributeFilename]:
"""Get the document attribute file name attribute in the document."""
return next(filter(lambda x: isinstance(x, DocumentAttributeFilename),
self.document.attributes), None)
@cached_property
def file_name(self) -> str:
"""Get the file name."""
return self.filename_attr.file_name if self.filename_attr else 'Unknown'
@property
def file_name_extension(self) -> str:
"""Get the file name extension."""
parts = self.file_name.rsplit(".", 1)
return parts[-1] if len(parts) >= 2 else ""
@property
def document(self):
"""Get the message document."""
return self.message.document
@property
def size(self) -> int:
"""Get the file size."""
return self.document.size
def __eq__(self, other: 'DownloadFile'):
"""Compare download files by their file name."""
return self.file_name == other.file_name
class DownloadSplitFilesBase:
"""Iterate over complete and split files. Base class to inherit."""
def __init__(self, messages: Iterable[Message]):
self.messages = messages
def get_iterator(self) -> Iterator[DownloadFile]:
"""Get an iterator with the download files."""
raise NotImplementedError
def __iter__(self) -> 'DownloadSplitFilesBase':
"""Set the iterator from the get_iterator method."""
self._iterator = self.get_iterator()
return self
def __next__(self) -> 'DownloadFile':
"""Get the next download file in the iterator."""
if self._iterator is None:
self._iterator = self.get_iterator()
return next(self._iterator)
class KeepDownloadSplitFiles(DownloadSplitFilesBase):
"""Download split files without join it."""
def get_iterator(self) -> Iterator[DownloadFile]:
"""Get an iterator with the download files."""
return map(lambda message: DownloadFile(message), self.messages)
class JoinDownloadSplitFiles(DownloadSplitFilesBase):
"""Download split files and join it."""
def get_iterator(self) -> Iterator[DownloadFile]:
"""Get an iterator with the download files. This method applies the join strategy and
joins the files after download it.
"""
current_join_strategy: Optional[JoinStrategyBase] = None
for message in self.messages:
download_file = DownloadFile(message)
yield download_file
if current_join_strategy and current_join_strategy.is_part(download_file):
# There is a bundle in process and the download file is part of it. Add the download
# file to the bundle.
current_join_strategy.add_download_file(download_file)
elif current_join_strategy and not current_join_strategy.is_part(download_file):
# There is a bundle in process and the download file is not part of it. Join the files
# in the bundle and finish it.
current_join_strategy.join_download_files()
current_join_strategy = None
if current_join_strategy is None:
# There is no bundle in process. Get the current bundle if the file has a strategy
# available.
current_join_strategy = get_join_strategy(download_file)
else:
# After finish all the files, join the latest bundle.
if current_join_strategy:
current_join_strategy.join_download_files()