txlyre пре 2 година
родитељ
комит
649da0148f
4 измењених фајлова са 105 додато и 8 уклоњено
  1. 17 2
      actions.py
  2. 54 5
      commands.py
  3. 7 1
      models.py
  4. 27 0
      utils.py

+ 17 - 2
actions.py

@@ -7,7 +7,7 @@ from telethon.tl.functions.messages import UploadMediaRequest, GetStickerSetRequ
 from telethon.tl.types import InputStickerSetID, InputStickerSetShortName, InputStickerSetItem, InputMediaUploadedDocument, InputPeerSelf
 from tortoise.expressions import F
 
-from models import Action, Gif, UserColor, StickerPack, Admin
+from models import Action, Gif, UserColor, StickerPack, Admin, BirthDay
 from utils import is_valid_name
 from config import config
 
@@ -126,4 +126,19 @@ async def add_sticker(bot, file, emoji):
     ))
     
   return get_input_document(pack.documents[-1])
-  
+
+async def get_birthdays(peer_id):
+  return await BirthDay.filter(peer_id=peer_id).all()
+
+async def add_or_update_birthday(peer_id, user, date):
+  birthday = await BirthDay.filter(peer_id=peer_id, user_id=user.id).first()
+  if birthday:
+    return False
+
+  await BirthDay(
+    peer_id=peer_id,
+    user_id=user.id,
+    date=date
+  ).save()
+
+  return True

+ 54 - 5
commands.py

@@ -1,10 +1,11 @@
 from sys import executable
 from asyncio import create_subprocess_shell
 from asyncio.subprocess import PIPE
+from datetime import datetime, timedelta
 
 from ujson import dumps
 from tortoise.exceptions import IntegrityError
-from telethon.utils import get_display_name
+from telethon.utils import get_display_name, get_peer_id
 from aiofiles.os import remove
 from emoji import is_emoji
 
@@ -16,9 +17,11 @@ from actions import (
   add_sticker,
   assign_color,
   add_admin,
-  delete_admin
+  delete_admin,
+  add_or_update_birthday,
+  get_birthdays
 )
-from utils import make_temporary_filename, parse_kind
+from utils import WORDS_TABLE,  make_temporary_filename, parse_kind, get_link_to_user, get_word_kind
 
 class Handler:
   def __init__(self, handler, is_restricted=False):
@@ -203,7 +206,51 @@ async def save_handler(bot, event, command):
       reply_to=message
     )
   finally:
-    await remove(path)    
+    await remove(path)
+
+async def bday_handler(bot, event, command):
+  if command.argc >= 1:
+    try:
+      date = datetime.strptime(' '.join(command.args), '%d.%m.%Y')
+    except ValueError:
+      await event.reply('Дата не может быть распознана. Пожалуйста, введите свой день рождения в следующем формате: 01.01.1970 (день, месяц, год).')
+
+      return
+
+    if await add_or_update_birthday(get_peer_id(event.peer_id), event.sender, date):
+      await event.reply('День рождения успешно добавлен!!!')
+    else:
+      await event.reply('День рождения успешно обновлён!!!')
+
+    return
+
+  birthdays = await get_birthdays(get_peer_id(event.peer_id))
+  if not birthdays:
+    await event.reply('Пока пусто…')
+
+    return
+
+  birthdays = sorted(birthdays, key=lambda birthday: birthday.date)
+
+  birthdays_list = ''
+  for birthday in birthdays:   
+    now = datetime.now().date()
+    birthday_date = birthday.date.replace(year=now.year)
+
+    if now > birthday_date:
+      birthday_date = birthday.date.replace(year=now.year + 1)
+
+    delta = birthday_date - now
+    age = (birthday_date - birthday.date).days // 365
+    age_now = age - 1
+
+    birthdays_list += '- '
+    birthdays_list += get_link_to_user(await bot.get_entity(birthday.user_id))
+    birthdays_list += ': '
+    birthdays_list += birthday.date.strftime('%d.%m.%Y')
+    birthdays_list += f' (через {delta.days} {WORDS_TABLE["день"][get_word_kind(delta.days)]} исполнится {age} {WORDS_TABLE["год"][get_word_kind(age)]}; сейчас {age_now} {WORDS_TABLE["год"][get_word_kind(age_now)]})\n'
+
+  await event.reply(f'Дни рождения:\n{birthdays_list}')
 
 COMMANDS = {
   'newadmin':  Handler(newadmin_handler,  is_restricted=True),
@@ -212,5 +259,7 @@ COMMANDS = {
   'delaction': Handler(delaction_handler, is_restricted=True),
   'addgif':    Handler(addgif_handler,    is_restricted=True),
 
-  'save':      Handler(save_handler)
+  'save':      Handler(save_handler),
+
+  'bday':      Handler(bday_handler)
 }

+ 7 - 1
models.py

@@ -1,5 +1,5 @@
 from tortoise.models import Model
-from tortoise.fields import IntField, BigIntField, CharField, TextField, BooleanField, ForeignKeyField
+from tortoise.fields import IntField, BigIntField, CharField, TextField, BooleanField, DateField, ForeignKeyField
 
 class Action(Model):
   id = IntField(pk=True)
@@ -28,3 +28,9 @@ class StickerPack(Model):
 class Admin(Model):
   id = IntField(pk=True)
   user_id = IntField(unique=True)
+
+class BirthDay(Model):
+  id = IntField(pk=True)
+  peer_id = BigIntField(unique=True)
+  user_id = IntField(unique=True)
+  date = DateField()

+ 27 - 0
utils.py

@@ -14,6 +14,16 @@ class Kind(IntEnum):
   CAN_APPLY_TO_SELF    = 1
   NO_TARGET            = 2
 
+class WordKind(IntEnum):
+  FIRST  = 0 # «год» / «день».
+  SECOND = 1 # «лет» / «дней».
+  THIRD  = 2 # «года» / «дня».
+
+WORDS_TABLE = {
+ 'год': ('год', 'лет', 'года'),
+ 'день': ('день', 'дней', 'дня')
+}
+
 def parse_command(text):
   text = text.strip()
 
@@ -71,3 +81,20 @@ def parse_kind(kind):
 
   return kind
 
+# Bruh?
+def get_word_kind(number):
+  if number == 0 or 5 <= number <= 20:
+    return WordKind.SECOND
+
+  last_digit = number % 10
+
+  if last_digit == 0:
+    return WordKind.SECOND
+
+  if last_digit == 1:
+    return WordKind.FIRST
+
+  if 5 <= last_digit <= 9:
+    return WordKind.SECOND
+
+  return WordKind.THIRD