openKriemy/utils.py

import re

from uuid import uuid4
from enum import IntEnum
from datetime import datetime
from ipaddress import ip_address
from collections import namedtuple

from aiofiles.os import mkdir
from aiofiles.os import path

from telethon.utils import get_display_name
from telethon.tl.types import (
  MessageEntityBold,
  MessageEntityItalic,
  MessageEntityStrike,
  MessageEntityCode,
  MessageEntityPre,
  MessageEntitySpoiler,
  MessageEntityUnderline,
  MessageEntityUrl,
  MessageEntityMention,
  MessageEntityBotCommand
)

Command = namedtuple(
  'Command', 
  'name argc args args_string'
)

Age = namedtuple(
  'Age',
  'days_until age age_now date_string'
)

ANSI_ESCAPE_RE = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')

class Kind(IntEnum):
  CANNOT_APPLY_TO_SELF = 0
  CAN_APPLY_TO_SELF    = 1
  NO_TARGET            = 2
  NO_TARGET_MAYBE      = 3

def parse_command(text):
  text = text.strip()

  if not text.startswith('/'):
    raise ValueError

  text = text.split(' ')

  if len(text) < 1:
    raise ValueError

  command = text[0][1:].lower()

  if '@' in command:
    command = command.split('@')
    command = command[0]

  args = text[1:]
  argc = len(args)
  args_string = ' '.join(args)

  return Command(
    name=command, 
    argc=argc, 
    args=args, 
    args_string=args_string
  )

def get_user_name(user):
  full_name = get_display_name(user)
  if not full_name:
    full_name = user.username
  
  if not full_name:
    full_name = '?'

  return full_name

def get_link_to_user(user):
  full_name = get_user_name(user)

  if user.username:
    return f'[{full_name}](@{user.username})'

  return f'[{full_name}](tg://user?id={user.id})'

def is_valid_name(name):
  return name.isidentifier()

def make_temporary_filename(ext):
  uid = uuid4().hex

  return f'tmp_{uid}.{ext}'
  
CACHE_DIR = './cache'

async def make_cache_filename(id, ext):
  if not await path.isdir(CACHE_DIR):
    await mkdir(CACHE_DIR)

  return f'{CACHE_DIR}/{id}.{ext}'

def parse_kind(kind):
  kind = int(kind)

  if kind < 0 or kind > 3:
    raise ValueError

  return kind

def calculate_age(date):
  now = datetime.now().date()
  birthday_date = date.replace(year=now.year)

  if now > birthday_date:
    birthday_date = date.replace(year=now.year + 1)

  delta = birthday_date - now
  age = (birthday_date - date).days // 365
  age_now = age - 1

  return Age(
    days_until=delta.days,
    age=age,
    age_now=age_now,
    date_string=date.strftime('%d.%m.%Y')
  )

DELIMITERS = {
  MessageEntityBold:       '**',
  MessageEntityItalic:     '__',
  MessageEntityStrike:     '~~',
  MessageEntityCode:       '`',
  MessageEntityPre:        '`',
  MessageEntityUnderline:  '\ue000',
  MessageEntitySpoiler:    '\ue001',
  MessageEntityUrl:        '\ue002',
  MessageEntityMention:    '\ue003',
  MessageEntityBotCommand: '\ue003'
}

class LookupTable:
  def __init__(self):
    self._start = {}
    self._end = {}

  def insert(self, entity):
    delimiter = DELIMITERS.get(type(entity))
    if not delimiter:
      return

    start = entity.offset
    end = entity.offset + entity.length

    if start in self._start:
      self._start[start].append(delimiter)
    else:
      self._start[start] = [delimiter]

    if end in self._end:
      self._end[end].insert(0, delimiter)
    else:
      self._end[end] = [delimiter]

  def _lookup(self, position):
    if position in self._start:
      return ''.join(self._start[position])
    elif position in self._end:
      return ''.join(self._end[position]) 

  def process(self, text):
    text = list(text) + ['']
    result = ''

    for position, character in zip(range(len(text)), text):     
      delimiter = self._lookup(position)
      if delimiter:
        result += delimiter

      result += character

    return result

def unparse(text, entities):
  table = LookupTable()

  for entity in entities:
    table.insert(entity)

  return table.process(text)

def is_valid_ip(ip):
  try:
    ip_address(ip)
  except:
    return False

  return True

def remove_ansi_escapes(text):
  return ANSI_ESCAPE_RE.sub('', text)