txlyre пре 3 година
родитељ
комит
4c55f4e426
5 измењених фајлова са 55 додато и 24 уклоњено
  1. 6 3
      actions.py
  2. 9 4
      commands.py
  3. 1 1
      models.py
  4. 23 16
      openkriemy.py
  5. 16 0
      utils.py

+ 6 - 3
actions.py

@@ -33,11 +33,11 @@ async def delete_admin(user):
 
   await admin.delete()
 
-async def create_action(name, template, can_apply_to_self):
+async def create_action(name, template, kind):
   if not is_valid_name(name):
     raise SyntaxError
 
-  await Action(name=name, template=template, can_apply_to_self=can_apply_to_self).save()
+  await Action(name=name, template=template, kind=kind).save()
    
 async def find_action(name):
   if not is_valid_name(name):
@@ -75,11 +75,14 @@ async def create_new_pack(bot, sticker):
   set_id = last_pack.id + 1 if last_pack else 1
  
   user = await bot.get_entity(config.USER)
+
+  me = await bot.get_me()
+  bot_username = me.username
   
   pack = await bot(CreateStickerSetRequest(
     user_id=user.id,
     title=f'Messages #{set_id}.',
-    short_name=f'messages{set_id}_by_openkriemybot',
+    short_name=f'messages{set_id}_by_{bot_username}',
     stickers=[sticker]
   ))
 

+ 9 - 4
commands.py

@@ -18,7 +18,7 @@ from actions import (
   add_admin,
   delete_admin
 )
-from utils import make_temporary_filename
+from utils import make_temporary_filename, parse_kind
 
 class Handler:
   def __init__(self, handler, is_restricted=False):
@@ -71,14 +71,19 @@ async def deladmin_handler(bot, event, command):
 
 async def newaction_handler(bot, event, command):
   if command.argc < 3:
-    await event.reply('Пожалуйста, укажите флаг (разрешить ли применять действие к самому себе), имя и шаблон действия!')
+    await event.reply('Пожалуйста, укажите тип, имя и шаблон действия!')
 
     return
 
-  can_apply_to_self = command.args[0] == '1'
+  try:
+    kind = parse_kind(command.args[0])
+  except ValueError:
+    await event.reply('Неверный тип действия!!!')
+
+    return
 
   try:
-    await create_action(command.args[1], ' '.join(command.args[2:]), can_apply_to_self)
+    await create_action(command.args[1], ' '.join(command.args[2:]), kind)
   except SyntaxError:
     await event.reply('Недопустимое имя действия!!!')
 

+ 1 - 1
models.py

@@ -5,7 +5,7 @@ class Action(Model):
   id = IntField(pk=True)
   name = CharField(max_length=64, unique=True)
   template = TextField()
-  can_apply_to_self = BooleanField(default=True)
+  kind = IntField()
 
 class Gif(Model):
   id = IntField(pk=True)

+ 23 - 16
openkriemy.py

@@ -4,7 +4,11 @@ from telethon import TelegramClient
 from telethon.events import NewMessage
 from telethon.utils import resolve_bot_file_id
 
-from utils import parse_command, get_link_to_user
+from utils import (
+  parse_command,
+  get_link_to_user,
+  Kind
+)
 from config import config
 from db import init_db
 from actions import (
@@ -47,24 +51,27 @@ async def on_message(event):
     return
 
   reply_to = None
-  target = await event.get_reply_message()
+  target = None
 
-  if not target:
-    try:
-      target = await bot.get_entity(command.args[0])
-    except (ValueError, IndexError):
-      await event.reply('Это действие нужно применить на кого-то!')
+  if action.kind != Kind.NO_TARGET:
+    target = await event.get_reply_message()
 
-      return 
-  else:
-    reply_to = target
-    target = target.sender
+    if not target:
+      try:
+        target = await bot.get_entity(command.args[0])
+      except (ValueError, IndexError):
+        await event.reply('Это действие нужно применить на кого-то!')
 
-  if target.id == event.sender.id\
- and not action.can_apply_to_self:
-    await event.reply('Данное действие нельзя применять к самому себе...')
+        return 
+    else:
+      reply_to = target
+      target = target.sender
 
-    return 
+    if action.kind == Kind.CANNOT_APPLY_TO_SELF\
+   and target.id == event.sender.id:
+      await event.reply('Данное действие нельзя применять к самому себе...')
+
+      return 
     
   try:
     await event.delete()
@@ -73,7 +80,7 @@ async def on_message(event):
     
   text = action.template.format(**{
     'initiator': get_link_to_user(event.sender),
-    'target': get_link_to_user(target)
+    'target': get_link_to_user(target) if target else '?'
   })
   
   gif = await get_random_gif(action)

+ 16 - 0
utils.py

@@ -1,5 +1,7 @@
 from uuid import uuid4
+from enum import IntEnum
 from collections import namedtuple
+
 from telethon.utils import get_display_name
 
 Command = namedtuple(
@@ -7,6 +9,11 @@ Command = namedtuple(
   'name argc args args_string'
 )
 
+class Kind(IntEnum):
+  CANNOT_APPLY_TO_SELF = 0
+  CAN_APPLY_TO_SELF    = 1
+  NO_TARGET            = 2
+
 def parse_command(text):
   text = text.strip()
 
@@ -55,3 +62,12 @@ def make_temporary_filename(ext):
   uid = uuid4().hex
 
   return f'tmp_{uid}.{ext}'
+
+def parse_kind(kind):
+  kind = int(kind)
+
+  if kind < 0 or kind > 2:
+    raise ValueError
+
+  return kind
+