add aggregator app

2019-01-28 00:20:09 +03:00
parent 8643fc42f9
commit b7e68abb03
29 changed files with 665 additions and 105 deletions
@@ -0,0 +1,96 @@
+import os
+from datetime import datetime
+from tempfile import TemporaryDirectory
+
+import pytz
+from django.conf import settings
+from django.db import models
+from pyrogram import Chat as PyrogramChat, Message as PyrogramMessage, ChatPhoto as PyrogramChatPhoto
+from pyrogram.api.types.chat_photo import ChatPhoto as MTProtoChatPhoto
+
+
+class AggregationSource(models.Model):
+    owner = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
+    title = models.CharField(max_length=64)
+    chat_id = models.IntegerField(db_index=True)
+
+    def __str__(self):
+        return self.title
+
+
+class Chat(models.Model):
+    chat_id = models.IntegerField(db_index=True)
+    title = models.TextField()
+    username = models.CharField(max_length=64, null=True, blank=True)
+    photo = models.ImageField(null=True, blank=True)
+    photo_id = models.CharField(max_length=64, null=True, blank=True)
+
+    @classmethod
+    def from_obj(cls, chat: PyrogramChat, client):
+        obj, _ = Chat.objects.update_or_create(
+            chat_id=chat.id,
+            defaults={
+                'title': chat.title or '{} {}'.format(chat.first_name, chat.last_name).rstrip(),
+                'username': chat.username,
+            }
+        )
+        if chat.photo is None:
+            if obj.photo_id is not None:
+                obj.photo = obj.photo_id = None
+                obj.save()
+        else:
+            photo_file_id = chat.photo.small_file_id
+            if photo_file_id != obj.photo_id:
+                with TemporaryDirectory() as d:
+                    path = client.download_media(photo_file_id, os.path.join(d, ''), block=True)
+                    with open(path, 'rb') as f:
+                        obj.photo.save(os.path.basename(path), f, save=True)
+                obj.photo_id = chat.photo.small_file_id
+                obj.save()
+        return obj
+
+    def __str__(self):
+        return '{} (chat_id="{}")'.format(self.title, self.chat_id)
+
+
+class MessageManager(models.Manager):
+    def active_messages(self):
+        return self.get_queryset().filter(deleted=False, replaced_by__isnull=True)
+
+
+class Message(models.Model):
+    chat = models.ForeignKey(Chat, on_delete=models.CASCADE, related_name='messages')
+    message_id = models.PositiveIntegerField(db_index=True)
+    text = models.TextField(blank=True)
+    date = models.DateTimeField()
+    edit_date = models.DateTimeField(null=True, blank=True)
+    deleted = models.BooleanField(default=False)
+    replaced_by = models.ForeignKey('self', on_delete=models.CASCADE, null=True, blank=True)
+
+    objects = MessageManager()
+
+    @classmethod
+    def from_obj(cls, message: PyrogramMessage, client):
+        tz = pytz.timezone('UTC')
+        chat = Chat.from_obj(message.chat, client)
+        try:
+            old = Message.objects.active_messages().get(chat=chat, message_id=message.message_id)
+        except Message.DoesNotExist:
+            old = None
+        obj = Message.objects.create(
+            chat=chat,
+            message_id=message.message_id,
+            text=message.text.html if message.text else '',
+            date=tz.localize(datetime.utcfromtimestamp(message.date)),
+            edit_date=tz.localize(datetime.utcfromtimestamp(message.edit_date)) if message.edit_date else None,
+        )
+        if old is not None:
+            old.replaced_by = obj
+            old.save()
+        return obj
+
+    def __str__(self):
+        return 'id="{}" (chat_id="{}")'.format(self.message_id, self.chat.chat_id)
+
+    class Meta:
+        ordering = 'message_id', 'edit_date',