From 99cc4d80b245c936fa5f6b5d8f37cf1adf8c3060 Mon Sep 17 00:00:00 2001
From: Logan Williams <logan@bellingcat.com>
Date: Thu, 4 May 2023 16:23:24 +0200
Subject: [PATCH] Cache screenname ID lookup

---
 cisticola/transformer/telegram_telethon.py | 34 ++++++++++++++--------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/cisticola/transformer/telegram_telethon.py b/cisticola/transformer/telegram_telethon.py
index fd7c941..e120d3d 100644
--- a/cisticola/transformer/telegram_telethon.py
+++ b/cisticola/transformer/telegram_telethon.py
@@ -30,6 +30,8 @@ class TelegramTelethonTransformer(Transformer):
 
     posts_cache = {}
 
+    get_screenname_cache = {}
+
     def can_handle(self, data: ScraperResult) -> bool:
         scraper = data.scraper.split(' ')
         if scraper[0] == "TelegramTelethonScraper":
@@ -52,18 +54,26 @@ class TelegramTelethonTransformer(Transformer):
         self.client.connect()
 
     def get_screenname_from_id(self, channel_id):
-        try:
-            data = self.client.get_entity(channel_id)
-            if isinstance(data, types.User):
-                return (data.username, str(data.first_name or "") + " " + str(data.last_name or ""), "")
-            else:
-                return (data.username, data.title, "")
-        except ChannelPrivateError:
-            return ("", "", "ChannelPrivateError")
-        except ChannelInvalidError:
-            return ("", "", "ChannelInvalidError")
-        except ValueError:
-            return ("", "", "ValueError")
+        if channel_id in self.get_screenname_cache:
+            return self.get_screenname_cache[channel_id]
+        else:
+            output = ("", "", None)
+            
+            try:
+                data = self.client.get_entity(channel_id)
+                if isinstance(data, types.User):
+                    output = (data.username, str(data.first_name or "") + " " + str(data.last_name or ""), "")
+                else:
+                    output = (data.username, data.title, "")
+            except ChannelPrivateError:
+                output = ("", "", "ChannelPrivateError")
+            except ChannelInvalidError:
+                output = ("", "", "ChannelInvalidError")
+            except ValueError:
+                output = ("", "", "ValueError")
+
+            self.get_screenname_cache[channel_id] = output
+            return output
 
     def get_name_from_web_interface(self, orig_screenname, id):
         url = "https://t.me/s/" + orig_screenname + "/" + str(id)