From 43d16474c24ab06fa91ce6aa3e27ec7e13b07ea0 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 6 Aug 2019 19:58:15 +0300
Subject: [PATCH] Improve logging and make get_dialogs use iterators more

---
 mautrix_telegram/abstract_user.py             | 18 ++++++++----------
 mautrix_telegram/bot.py                       |  2 +-
 mautrix_telegram/puppet.py                    |  2 ++
 mautrix_telegram/user.py                      |  6 +++++-
 mautrix_telegram/web/provisioning/__init__.py |  3 +--
 5 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/mautrix_telegram/abstract_user.py b/mautrix_telegram/abstract_user.py
index 0acc0aed..101af930 100644
--- a/mautrix_telegram/abstract_user.py
+++ b/mautrix_telegram/abstract_user.py
@@ -13,7 +13,7 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
-from typing import Tuple, Optional, List, Union, Dict, TYPE_CHECKING
+from typing import Tuple, Optional, AsyncIterable, Union, Dict, TYPE_CHECKING
 from abc import ABC, abstractmethod
 import asyncio
 import logging
@@ -189,14 +189,12 @@ class AbstractUser(ABC):
         if UPDATE_TIME:
             UPDATE_TIME.labels(update_type=type(update).__name__).observe(time.time() - start_time)
 
-    async def get_dialogs(self, limit: int = None) -> List[Union[Chat, Channel]]:
-        if self.is_bot:
-            return []
-        dialogs = await self.client.get_dialogs(limit=limit)
-        return [dialog.entity for dialog in dialogs if (
-            not isinstance(dialog.entity, (User, ChatForbidden, ChannelForbidden))
-            and not (isinstance(dialog.entity, Chat)
-                     and (dialog.entity.deactivated or dialog.entity.left)))]
+    def get_dialogs(self, limit: int = None) -> AsyncIterable[Union[User, Chat, Channel]]:
+        return (dialog.entity async for dialog in
+                self.client.iter_dialogs(limit=limit, ignore_migrated=True, archived=False)
+                if isinstance(dialog.entity, (ChatForbidden, ChannelForbidden))
+                or (isinstance(dialog.entity, Chat)
+                    and (dialog.entity.deactivated or dialog.entity.left)))
 
     @property
     @abstractmethod
@@ -216,7 +214,7 @@ class AbstractUser(ABC):
         if not self.client:
             self._init_client()
         await self.client.connect()
-        self.log.debug("%s connected: %s", self.mxid, self.connected)
+        self.log.debug(f"{self.mxid if not self.is_bot else 'Bot'} connected: {self.connected}")
         return self
 
     async def ensure_started(self, even_if_no_session=False) -> 'AbstractUser':
diff --git a/mautrix_telegram/bot.py b/mautrix_telegram/bot.py
index b8cf19e8..f03d072a 100644
--- a/mautrix_telegram/bot.py
+++ b/mautrix_telegram/bot.py
@@ -43,7 +43,7 @@ ReplyFunc = Callable[[str], Awaitable[Message]]
 
 
 class Bot(AbstractUser):
-    log: logging.Logger = logging.getLogger("mau.bot")
+    log: logging.Logger = logging.getLogger("mau.user.bot")
     mxid_regex: Pattern = re.compile("@.+:.+")
 
     token: str
diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py
index 0eea73a2..a1ae68dc 100644
--- a/mautrix_telegram/puppet.py
+++ b/mautrix_telegram/puppet.py
@@ -108,6 +108,8 @@ class Puppet(CustomPuppetMixin):
         if self.custom_mxid:
             self.by_custom_mxid[self.custom_mxid] = self
 
+        self.log = self.log.getChild(str(self.id))
+
     @property
     def tgid(self) -> TelegramID:
         return self.id
diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py
index 7b765335..b20629ad 100644
--- a/mautrix_telegram/user.py
+++ b/mautrix_telegram/user.py
@@ -87,6 +87,8 @@ class User(AbstractUser):
         if tgid:
             self.by_tgid[tgid] = self
 
+        self.log = self.log.getChild(self.mxid)
+
     @property
     def name(self) -> str:
         return self.mxid
@@ -305,8 +307,10 @@ class User(AbstractUser):
         return await self._search_remote(query), True
 
     async def sync_dialogs(self, synchronous_create: bool = False) -> None:
+        if self.is_bot:
+            return
         creators = []
-        for entity in await self.get_dialogs(limit=config["bridge.sync_dialog_limit"] or None):
+        async for entity in self.get_dialogs(limit=config["bridge.sync_dialog_limit"] or None):
             portal = po.Portal.get_by_entity(entity)
             self.portals[portal.tgid_full] = portal
             creators.append(
diff --git a/mautrix_telegram/web/provisioning/__init__.py b/mautrix_telegram/web/provisioning/__init__.py
index 9487e8ce..3a81c19c 100644
--- a/mautrix_telegram/web/provisioning/__init__.py
+++ b/mautrix_telegram/web/provisioning/__init__.py
@@ -316,11 +316,10 @@ class ProvisioningAPI(AuthAPI):
             return err
 
         if not user.is_bot:
-            chats = await user.get_dialogs()
             return web.json_response([{
                 "id": get_peer_id(chat),
                 "title": chat.title,
-            } for chat in chats])
+            } async for chat in user.get_dialogs()])
         else:
             return web.json_response([{
                 "id": get_peer_id(chat.peer),