From 5ec984dd820a00d005306b6144b86b9b54346c99 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Wed, 7 Feb 2018 10:41:58 +0100
Subject: [PATCH] Allow adding events with the client.on decorator

---
 telethon/events/__init__.py | 14 ++++++++++++++
 telethon/telegram_client.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)
 create mode 100644 telethon/events/__init__.py

diff --git a/telethon/events/__init__.py b/telethon/events/__init__.py
new file mode 100644
index 00000000..d4d1e019
--- /dev/null
+++ b/telethon/events/__init__.py
@@ -0,0 +1,14 @@
+import abc
+from ..tl import types, functions
+from ..extensions import markdown
+from .. import utils
+
+
+class _EventBuilder(abc.ABC):
+    @abc.abstractmethod
+    def build(self, update):
+        """Builds an event for the given update if possible, or returns None"""
+
+    @abc.abstractmethod
+    def resolve(self, client):
+        """Helper method to allow event builders to be resolved before usage"""
diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 9dda06af..c4f5d722 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -160,6 +160,8 @@ class TelegramClient(TelegramBareClient):
             **kwargs
         )
 
+        self._event_builders = []
+
         # Some fields to easy signing in. Let {phone: hash} be
         # a dictionary because the user may change their mind.
         self._phone_code_hash = {}
@@ -1623,6 +1625,41 @@ class TelegramClient(TelegramBareClient):
 
     # endregion
 
+    # region Event handling
+
+    def on(self, event):
+        """
+
+        Turns the given entity into a valid Telegram user or chat.
+
+        Args:
+            event (:obj:`_EventBuilder` | :obj:`type`):
+                The event builder class or instance to be used,
+                for instance ``events.NewMessage``.
+        """
+        if isinstance(event, type):
+            event = event()
+
+        event.resolve(self)
+
+        def decorator(f):
+            self._event_builders.append((event, f))
+            return f
+
+        if self._on_handler not in self.updates.handlers:
+            self.add_update_handler(self._on_handler)
+
+        return decorator
+
+    def _on_handler(self, update):
+        for builder, callback in self._event_builders:
+            event = builder.build(update)
+            if event:
+                event._client = self
+                callback(event)
+
+    # endregion
+
     # region Small utilities to make users' life easier
 
     def get_entity(self, entity):