Simplify event building logic

This will also fix some bugs where the event type being accessed
would not be available, since it is now built on-demand, without
the need to keep track for the count of each event type.

telethon/client/telegrambaseclient.py

@@ -274,11 +274,6 @@ class TelegramBaseClient(abc.ABC):
         self._event_resolve_lock = asyncio.Lock()
         self._conversations = {}
 
-        # Keep track of how many event builders there are for
-        # each type {type: count}. If there's at least one then
-        # the event will be built, and the same event be reused.
-        self._event_builders_count = collections.defaultdict(int)
-
         # Default parse mode
         self._parse_mode = markdown
 

telethon/client/updates.py

@@ -90,7 +90,6 @@ class UpdateMethods(UserMethods):
             event = events.Raw()
 
         self._events_pending_resolve.append(event)
-        self._event_builders_count[type(event)] += 1
         self._event_builders.append((event, callback))
 
     def remove_event_handler(self, callback, event=None):
@@ -109,11 +108,6 @@ class UpdateMethods(UserMethods):
             i -= 1
             ev, cb = self._event_builders[i]
             if cb == callback and (not event or isinstance(ev, event)):
-                type_ev = type(ev)
-                self._event_builders_count[type_ev] -= 1
-                if not self._event_builders_count[type_ev]:
-                    del self._event_builders_count[type_ev]
-
                 del self._event_builders[i]
                 found += 1
 
@@ -266,44 +260,23 @@ class UpdateMethods(UserMethods):
 
             self._events_pending_resolve.clear()
 
-        # TODO We can improve this further
-        # If we had a way to get all event builders for
-        # a type instead looping over them all always.
-        built = {builder: builder.build(update)
-                 for builder in self._event_builders_count}
-
+        built = EventBuilderDict(self, update)
         if self._conversations:
-            for ev_type in (events.NewMessage, events.MessageEdited,
-                            events.MessageRead):
-                if ev_type not in built:
-                    built[ev_type] = ev_type.build(update)
-
             for conv in self._conversations.values():
-                # TODO Cleaner way to do this? Maybe just pass built always?
                 if built[events.NewMessage]:
-                    built[events.NewMessage]._set_client(self)
                     conv._on_new_message(built[events.NewMessage])
                 if built[events.MessageEdited]:
-                    built[events.MessageEdited]._set_client(self)
                     conv._on_edit(built[events.MessageEdited])
                 if built[events.MessageRead]:
-                    built[events.MessageRead]._set_client(self)
                     conv._on_read(built[events.MessageRead])
-
                 if conv._custom:
-                    await conv._check_custom(built, update)
+                    await conv._check_custom(built)
 
         for builder, callback in self._event_builders:
             event = built[type(builder)]
             if not event or not builder.filter(event):
                 continue
 
-            if hasattr(event, '_set_client'):
-                event._set_client(self)
-            else:
-                event._client = self
-
-            event.original_update = update
             try:
                 await callback(event)
             except events.StopPropagation:
@@ -328,3 +301,26 @@ class UpdateMethods(UserMethods):
             __log__.info('Failed to get current state: %r', e)
 
     # endregion
+
+
+class EventBuilderDict:
+    """
+    Helper "dictionary" to return events from types and cache them.
+    """
+    def __init__(self, client, update):
+        self.client = client
+        self.update = update
+
+    def __getitem__(self, builder):
+        try:
+            return self.__dict__[builder]
+        except KeyError:
+            event = self.__dict__[builder] = builder.build(self.update)
+            if event:
+                event.original_update = self.update
+                if hasattr(event, '_set_client'):
+                    event._set_client(self.client)
+                else:
+                    event._client = self.client
+
+            return event

telethon/tl/custom/conversation.py

@@ -273,26 +273,18 @@ class Conversation(ChatGetter):
         self._custom[counter] = (event, future, False)
         return await result()
 
-    async def _check_custom(self, built, update):
+    async def _check_custom(self, built):
         # TODO This code is quite much a copy paste of registering events
         # in the client, resolving them and setting the client; perhaps
         # there is a better way?
         for i, (ev, fut, resolved) in self._custom.items():
             ev_type = type(ev)
-            if ev_type not in built:
-                built[ev_type] = ev.build(update)
-
             if built[ev_type]:
                 if not resolved:
                     await ev.resolve(self._client)
                     self._custom[i] = (ev, fut, True)
 
                 if ev.filter(built[ev_type]):
-                    if hasattr(built[ev_type], '_set_client'):
-                        built[ev_type]._set_client(self._client)
-                    else:
-                        built[ev_type]._client = self._client
-
                     fut.set_result(built[ev_type])
 
     def _on_new_message(self, response):