Fix direct mutation of objects in friendly methods

This commit is contained in:
Lonami Exo 2022-02-07 09:28:39 +01:00
parent ae43a47602
commit d87b68a756
6 changed files with 65 additions and 61 deletions

View File

@ -3,6 +3,7 @@ import inspect
import itertools import itertools
import string import string
import typing import typing
import dataclasses
from .. import errors, _tl from .. import errors, _tl
from .._misc import helpers, utils, requestiter, tlobject, enums, hints from .._misc import helpers, utils, requestiter, tlobject, enums, hints
@ -19,11 +20,9 @@ _MAX_PROFILE_PHOTO_CHUNK_SIZE = 100
class _ChatAction: class _ChatAction:
def __init__(self, client, chat, action, *, delay, auto_cancel): def __init__(self, client, chat, action, *, delay, auto_cancel):
self._client = client self._client = client
self._chat = chat
self._action = action
self._delay = delay self._delay = delay
self._auto_cancel = auto_cancel self._auto_cancel = auto_cancel
self._request = None self._request = _tl.fn.messages.SetTyping(chat, action)
self._task = None self._task = None
self._running = False self._running = False
@ -31,14 +30,7 @@ class _ChatAction:
return self._once().__await__() return self._once().__await__()
async def __aenter__(self): async def __aenter__(self):
self._chat = await self._client.get_input_entity(self._chat) self._request = dataclasses.replace(self._request, peer=await self._client.get_input_entity(self._request.peer))
# Since `self._action` is passed by reference we can avoid
# recreating the request all the time and still modify
# `self._action.progress` directly in `progress`.
self._request = _tl.fn.messages.SetTyping(
self._chat, self._action)
self._running = True self._running = True
self._task = asyncio.create_task(self._update()) self._task = asyncio.create_task(self._update())
return self return self
@ -55,7 +47,7 @@ class _ChatAction:
self._task = None self._task = None
async def _once(self): async def _once(self):
self._chat = await self._client.get_input_entity(self._chat) self._request = dataclasses.replace(self._request, peer=await self._client.get_input_entity(self._request.peer))
await self._client(_tl.fn.messages.SetTyping(self._chat, self._action)) await self._client(_tl.fn.messages.SetTyping(self._chat, self._action))
async def _update(self): async def _update(self):
@ -93,8 +85,11 @@ class _ChatAction:
}[enums.Action(action)] }[enums.Action(action)]
def progress(self, current, total): def progress(self, current, total):
if hasattr(self._action, 'progress'): if hasattr(self._request.action, 'progress'):
self._action.progress = 100 * round(current / total) self._request = dataclasses.replace(
self._request,
action=dataclasses.replace(self._request.action, progress=100 * round(current / total))
)
class _ParticipantsIter(requestiter.RequestIter): class _ParticipantsIter(requestiter.RequestIter):
@ -190,8 +185,8 @@ class _ParticipantsIter(requestiter.RequestIter):
# Most people won't care about getting exactly 12,345 # Most people won't care about getting exactly 12,345
# members so it doesn't really matter not to be 100% # members so it doesn't really matter not to be 100%
# precise with being out of the offset/limit here. # precise with being out of the offset/limit here.
self.request.limit = min( self.request = dataclasses.replace(self.request, limit=min(
self.limit - self.request.offset, _MAX_PARTICIPANTS_CHUNK_SIZE) self.limit - self.request.offset, _MAX_PARTICIPANTS_CHUNK_SIZE))
if self.request.offset > self.limit: if self.request.offset > self.limit:
return True return True
@ -199,7 +194,7 @@ class _ParticipantsIter(requestiter.RequestIter):
participants = await self.client(self.request) participants = await self.client(self.request)
self.total = participants.count self.total = participants.count
self.request.offset += len(participants.participants) self.request = dataclasses.replace(self.request, offset=self.request.offset + len(participants.participants))
users = {user.id: user for user in participants.users} users = {user.id: user for user in participants.users}
for participant in participants.participants: for participant in participants.participants:
if isinstance(participant, _tl.ChannelParticipantBanned): if isinstance(participant, _tl.ChannelParticipantBanned):
@ -253,20 +248,20 @@ class _AdminLogIter(requestiter.RequestIter):
) )
async def _load_next_chunk(self): async def _load_next_chunk(self):
self.request.limit = min(self.left, _MAX_ADMIN_LOG_CHUNK_SIZE) self.request = dataclasses.replace(self.request, limit=min(self.left, _MAX_ADMIN_LOG_CHUNK_SIZE))
r = await self.client(self.request) r = await self.client(self.request)
entities = {utils.get_peer_id(x): x entities = {utils.get_peer_id(x): x
for x in itertools.chain(r.users, r.chats)} for x in itertools.chain(r.users, r.chats)}
self.request.max_id = min((e.id for e in r.events), default=0) self.request = dataclasses.replace(self.request, max_id=min((e.id for e in r.events), default=0))
for ev in r.events: for ev in r.events:
if isinstance(ev.action, if isinstance(ev.action,
_tl.ChannelAdminLogEventActionEditMessage): _tl.ChannelAdminLogEventActionEditMessage):
ev.action.prev_message = _custom.Message._new( ev = dataclasses.replace(ev, action=dataclasses.replace(
self.client, ev.action.prev_message, entities, self.entity) ev.action,
prev_message=_custom.Message._new(self.client, ev.action.prev_message, entities, self.entity),
ev.action.new_message = _custom.Message._new( new_message=_custom.Message._new(self.client, ev.action.new_message, entities, self.entity)
self.client, ev.action.new_message, entities, self.entity) ))
elif isinstance(ev.action, elif isinstance(ev.action,
_tl.ChannelAdminLogEventActionDeleteMessage): _tl.ChannelAdminLogEventActionDeleteMessage):
@ -308,7 +303,7 @@ class _ProfilePhotoIter(requestiter.RequestIter):
) )
if self.limit == 0: if self.limit == 0:
self.request.limit = 1 self.request = dataclasses.replace(self.request, limit=1)
result = await self.client(self.request) result = await self.client(self.request)
if isinstance(result, _tl.photos.Photos): if isinstance(result, _tl.photos.Photos):
self.total = len(result.photos) self.total = len(result.photos)
@ -319,7 +314,7 @@ class _ProfilePhotoIter(requestiter.RequestIter):
self.total = getattr(result, 'count', None) self.total = getattr(result, 'count', None)
async def _load_next_chunk(self): async def _load_next_chunk(self):
self.request.limit = min(self.left, _MAX_PROFILE_PHOTO_CHUNK_SIZE) self.request = dataclasses.replace(self.request, limit=min(self.left, _MAX_PROFILE_PHOTO_CHUNK_SIZE))
result = await self.client(self.request) result = await self.client(self.request)
if isinstance(result, _tl.photos.Photos): if isinstance(result, _tl.photos.Photos):
@ -338,7 +333,7 @@ class _ProfilePhotoIter(requestiter.RequestIter):
if len(self.buffer) < self.request.limit: if len(self.buffer) < self.request.limit:
self.left = len(self.buffer) self.left = len(self.buffer)
else: else:
self.request.offset += len(result.photos) self.request = dataclasses.replace(self.request, offset=self.request.offset + len(result.photos))
else: else:
# Some broadcast channels have a photo that this request doesn't # Some broadcast channels have a photo that this request doesn't
# retrieve for whatever random reason the Telegram server feels. # retrieve for whatever random reason the Telegram server feels.
@ -368,8 +363,11 @@ class _ProfilePhotoIter(requestiter.RequestIter):
if len(result.messages) < self.request.limit: if len(result.messages) < self.request.limit:
self.left = len(self.buffer) self.left = len(self.buffer)
elif result.messages: elif result.messages:
self.request.add_offset = 0 self.request = dataclasses.replace(
self.request.offset_id = result.messages[-1].id self.request,
add_offset=0,
offset_id=result.messages[-1].id
)
def get_participants( def get_participants(

View File

@ -2,6 +2,7 @@ import asyncio
import inspect import inspect
import itertools import itertools
import typing import typing
import dataclasses
from .. import errors, _tl from .. import errors, _tl
from .._misc import helpers, utils, requestiter, hints from .._misc import helpers, utils, requestiter, hints
@ -49,7 +50,7 @@ class _DialogsIter(requestiter.RequestIter):
self.ignore_migrated = ignore_migrated self.ignore_migrated = ignore_migrated
async def _load_next_chunk(self): async def _load_next_chunk(self):
self.request.limit = min(self.left, _MAX_CHUNK_SIZE) self.request = dataclasses.replace(self.request, limit=min(self.left, _MAX_CHUNK_SIZE))
r = await self.client(self.request) r = await self.client(self.request)
self.total = getattr(r, 'count', len(r.dialogs)) self.total = getattr(r, 'count', len(r.dialogs))
@ -103,10 +104,13 @@ class _DialogsIter(requestiter.RequestIter):
for d in reversed(r.dialogs) for d in reversed(r.dialogs)
)), None) )), None)
self.request.exclude_pinned = True self.request = dataclasses.replace(
self.request.offset_id = last_message.id if last_message else 0 self.request,
self.request.offset_date = last_message.date if last_message else None exclude_pinned=True,
self.request.offset_peer = self.buffer[-1].input_entity offset_id=last_message.id if last_message else 0,
offset_date=last_message.date if last_message else None,
offset_peer=self.buffer[-1].input_entity,
)
class _DraftsIter(requestiter.RequestIter): class _DraftsIter(requestiter.RequestIter):

View File

@ -5,6 +5,7 @@ import pathlib
import typing import typing
import inspect import inspect
import asyncio import asyncio
import dataclasses
from .._crypto import AES from .._crypto import AES
from .._misc import utils, helpers, requestiter, tlobject, hints, enums from .._misc import utils, helpers, requestiter, tlobject, hints, enums
@ -55,7 +56,7 @@ class _DirectDownloadIter(requestiter.RequestIter):
self.left = len(self.buffer) self.left = len(self.buffer)
await self.close() await self.close()
else: else:
self.request.offset += self._stride self.request = dataclasses.replace(self.request, offset=self.request.offset + self._stride)
async def _request(self): async def _request(self):
try: try:
@ -102,7 +103,7 @@ class _DirectDownloadIter(requestiter.RequestIter):
if document.id != self.request.location.id: if document.id != self.request.location.id:
raise raise
self.request.location.file_reference = document.file_reference self.request.location = dataclasses.replace(self.request.location, file_reference=document.file_reference)
return await self._request() return await self._request()
async def close(self): async def close(self):
@ -134,18 +135,18 @@ class _GenericDownloadIter(_DirectDownloadIter):
before = self.request.offset before = self.request.offset
# 1.2. We have to fetch from a valid offset, so remove that bad part # 1.2. We have to fetch from a valid offset, so remove that bad part
self.request.offset -= bad self.request = dataclasses.replace(self.request, offset=self.request.offset - bad)
done = False done = False
while not done and len(data) - bad < self._chunk_size: while not done and len(data) - bad < self._chunk_size:
cur = await self._request() cur = await self._request()
self.request.offset += self.request.limit self.request = dataclasses.replace(self.request, offset=self.request.offset - self.request.limit)
data += cur data += cur
done = len(cur) < self.request.limit done = len(cur) < self.request.limit
# 1.3 Restore our last desired offset # 1.3 Restore our last desired offset
self.request.offset = before self.request = dataclasses.replace(self.request, offset=before)
# 2. Fill the buffer with the data we have # 2. Fill the buffer with the data we have
# 2.1. Slicing `bytes` is expensive, yield `memoryview` instead # 2.1. Slicing `bytes` is expensive, yield `memoryview` instead
@ -157,7 +158,7 @@ class _GenericDownloadIter(_DirectDownloadIter):
self.buffer.append(mem[i:i + self._chunk_size]) self.buffer.append(mem[i:i + self._chunk_size])
# 2.3. We will yield this offset, so move to the next one # 2.3. We will yield this offset, so move to the next one
self.request.offset += self._stride self.request = dataclasses.replace(self.request, offset=self.request.offset + self._stride)
# 2.4. If we are in the last chunk, we will return the last partial data # 2.4. If we are in the last chunk, we will return the last partial data
if done: if done:
@ -172,7 +173,7 @@ class _GenericDownloadIter(_DirectDownloadIter):
# 3. Be careful with the offsets. Re-fetching a bit of data # 3. Be careful with the offsets. Re-fetching a bit of data
# is fine, since it greatly simplifies things. # is fine, since it greatly simplifies things.
# TODO Try to not re-fetch data # TODO Try to not re-fetch data
self.request.offset -= self._stride self.request = dataclasses.replace(self.request, offset=self.request.offset - self._stride)
async def download_profile_photo( async def download_profile_photo(

View File

@ -87,7 +87,7 @@ def _get_response_message(self: 'TelegramClient', request, result, input_chat):
elif isinstance(update, ( elif isinstance(update, (
_tl.UpdateNewChannelMessage, _tl.UpdateNewMessage)): _tl.UpdateNewChannelMessage, _tl.UpdateNewMessage)):
update.message = _custom.Message._new(self, update.message, entities, input_chat) message = _custom.Message._new(self, update.message, entities, input_chat)
# Pinning a message with `updatePinnedMessage` seems to # Pinning a message with `updatePinnedMessage` seems to
# always produce a service message we can't map so return # always produce a service message we can't map so return
@ -97,20 +97,20 @@ def _get_response_message(self: 'TelegramClient', request, result, input_chat):
# #
# TODO this method is getting messier and messier as time goes on # TODO this method is getting messier and messier as time goes on
if hasattr(request, 'random_id') or utils.is_list_like(request): if hasattr(request, 'random_id') or utils.is_list_like(request):
id_to_message[update.message.id] = update.message id_to_message[message.id] = message
else: else:
return update.message return message
elif (isinstance(update, _tl.UpdateEditMessage) elif (isinstance(update, _tl.UpdateEditMessage)
and helpers._entity_type(request.peer) != helpers._EntityType.CHANNEL): and helpers._entity_type(request.peer) != helpers._EntityType.CHANNEL):
update.message = _custom.Message._new(self, update.message, entities, input_chat) message = _custom.Message._new(self, update.message, entities, input_chat)
# Live locations use `sendMedia` but Telegram responds with # Live locations use `sendMedia` but Telegram responds with
# `updateEditMessage`, which means we won't have `id` field. # `updateEditMessage`, which means we won't have `id` field.
if hasattr(request, 'random_id'): if hasattr(request, 'random_id'):
id_to_message[update.message.id] = update.message id_to_message[message.id] = message
elif request.id == update.message.id: elif request.id == message.id:
return update.message return message
elif (isinstance(update, _tl.UpdateEditChannelMessage) elif (isinstance(update, _tl.UpdateEditChannelMessage)
and utils.get_peer_id(request.peer) == and utils.get_peer_id(request.peer) ==

View File

@ -3,6 +3,7 @@ import itertools
import time import time
import typing import typing
import warnings import warnings
import dataclasses
from .._misc import helpers, utils, requestiter, hints from .._misc import helpers, utils, requestiter, hints
from ..types import _custom from ..types import _custom
@ -142,7 +143,7 @@ class _MessagesIter(requestiter.RequestIter):
and offset_date and not search and not offset_id: and offset_date and not search and not offset_id:
async for m in self.client.get_messages( async for m in self.client.get_messages(
self.entity, 1, offset_date=offset_date): self.entity, 1, offset_date=offset_date):
self.request.offset_id = m.id + 1 self.request = dataclasses.replace(self.request, offset_id=m.id + 1)
else: else:
self.request = _tl.fn.messages.GetHistory( self.request = _tl.fn.messages.GetHistory(
peer=self.entity, peer=self.entity,
@ -178,10 +179,10 @@ class _MessagesIter(requestiter.RequestIter):
self.last_id = 0 if self.reverse else float('inf') self.last_id = 0 if self.reverse else float('inf')
async def _load_next_chunk(self): async def _load_next_chunk(self):
self.request.limit = min(self.left, _MAX_CHUNK_SIZE) self.request = dataclasses.replace(self.request, limit=min(self.left, _MAX_CHUNK_SIZE))
if self.reverse and self.request.limit != _MAX_CHUNK_SIZE: if self.reverse and self.request.limit != _MAX_CHUNK_SIZE:
# Remember that we need -limit when going in reverse # Remember that we need -limit when going in reverse
self.request.add_offset = self.add_offset - self.request.limit self.request = dataclasses.replace(self.request, add_offset=self.add_offset - self.request.limit)
r = await self.client(self.request) r = await self.client(self.request)
self.total = getattr(r, 'count', len(r.messages)) self.total = getattr(r, 'count', len(r.messages))
@ -241,28 +242,28 @@ class _MessagesIter(requestiter.RequestIter):
""" """
After making the request, update its offset with the last message. After making the request, update its offset with the last message.
""" """
self.request.offset_id = last_message.id self.request = dataclasses.replace(self.request, offset_id=last_message.id)
if self.reverse: if self.reverse:
# We want to skip the one we already have # We want to skip the one we already have
self.request.offset_id += 1 self.request = dataclasses.replace(self.request, offset_id=self.request.offset_id + 1)
if isinstance(self.request, _tl.fn.messages.Search): if isinstance(self.request, _tl.fn.messages.Search):
# Unlike getHistory and searchGlobal that use *offset* date, # Unlike getHistory and searchGlobal that use *offset* date,
# this is *max* date. This means that doing a search in reverse # this is *max* date. This means that doing a search in reverse
# will break it. Since it's not really needed once we're going # will break it. Since it's not really needed once we're going
# (only for the first request), it's safe to just clear it off. # (only for the first request), it's safe to just clear it off.
self.request.max_date = None self.request = dataclasses.replace(self.request, max_date=None)
else: else:
# getHistory, searchGlobal and getReplies call it offset_date # getHistory, searchGlobal and getReplies call it offset_date
self.request.offset_date = last_message.date self.request = dataclasses.replace(self.request, offset_date=last_message.date)
if isinstance(self.request, _tl.fn.messages.SearchGlobal): if isinstance(self.request, _tl.fn.messages.SearchGlobal):
if last_message.input_chat: if last_message.input_chat:
self.request.offset_peer = last_message.input_chat self.request = dataclasses.replace(self.request, offset_peer=last_message.input_chat)
else: else:
self.request.offset_peer = _tl.InputPeerEmpty() self.request = dataclasses.replace(self.request, offset_peer=_tl.InputPeerEmpty())
self.request.offset_rate = getattr(response, 'next_rate', 0) self.request = dataclasses.replace(self.request, offset_rate=getattr(response, 'next_rate', 0))
class _IDsIter(requestiter.RequestIter): class _IDsIter(requestiter.RequestIter):

View File

@ -3,6 +3,7 @@ import datetime
import itertools import itertools
import time import time
import typing import typing
import dataclasses
from ..errors._custom import MultiError from ..errors._custom import MultiError
from ..errors._rpcbase import RpcError, ServerError, FloodError, InvalidDcError, UnauthorizedError from ..errors._rpcbase import RpcError, ServerError, FloodError, InvalidDcError, UnauthorizedError
@ -370,8 +371,7 @@ async def _get_input_dialog(self: 'TelegramClient', dialog):
""" """
try: try:
if dialog.SUBCLASS_OF_ID == 0xa21c9795: # crc32(b'InputDialogPeer') if dialog.SUBCLASS_OF_ID == 0xa21c9795: # crc32(b'InputDialogPeer')
dialog.peer = await self.get_input_entity(dialog.peer) return dataclasses.replace(dialog, peer=await self.get_input_entity(dialog.peer))
return dialog
elif dialog.SUBCLASS_OF_ID == 0xc91c90b6: # crc32(b'InputPeer') elif dialog.SUBCLASS_OF_ID == 0xc91c90b6: # crc32(b'InputPeer')
return _tl.InputDialogPeer(dialog) return _tl.InputDialogPeer(dialog)
except AttributeError: except AttributeError:
@ -388,7 +388,7 @@ async def _get_input_notify(self: 'TelegramClient', notify):
try: try:
if notify.SUBCLASS_OF_ID == 0x58981615: if notify.SUBCLASS_OF_ID == 0x58981615:
if isinstance(notify, _tl.InputNotifyPeer): if isinstance(notify, _tl.InputNotifyPeer):
notify.peer = await self.get_input_entity(notify.peer) return dataclasses.replace(notify, peer=await self.get_input_entity(notify.peer))
return notify return notify
except AttributeError: except AttributeError:
pass pass