python-app-rocketbot/rocketbot/utils.py

"""
General utility functions.
"""
import re
import sys
import traceback
from datetime import datetime, timedelta
from typing import Any, Optional, Type, Union

from discord import Guild
from discord.ext.commands import Cog, Group

def dump_stacktrace(e: Exception) -> None:
	print(e, file=sys.stderr)
	traceback.print_exception(type(e), e, e.__traceback__)

def timedelta_from_str(s: str) -> timedelta:
	"""
	Parses a timespan. Format examples:
	"30m"
	"10s"
	"90d"
	"1h30m"
	"73d18h22m52s"
	"""
	p: re.Pattern = re.compile('^(?:[0-9]+[dhms])+$')
	if p.match(s) is None:
		raise ValueError("Illegal timespan value '{s}'.")
	p = re.compile('([0-9]+)([dhms])')
	days: int = 0
	hours: int = 0
	minutes: int = 0
	seconds: int = 0
	for m in p.finditer(s):
		scalar = int(m.group(1))
		unit = m.group(2)
		if unit == 'd':
			days = scalar
		elif unit == 'h':
			hours = scalar
		elif unit == 'm':
			minutes = scalar
		elif unit == 's':
			seconds = scalar
	return timedelta(days=days, hours=hours, minutes=minutes, seconds=seconds)

def str_from_timedelta(td: timedelta) -> str:
	"""
	Encodes a timedelta as a str. E.g. "3d2h"
	"""
	d: int = td.days
	h: int = td.seconds // 3600
	m: int = (td.seconds // 60) % 60
	s: int = td.seconds % 60
	components: list[str] = []
	if d != 0:
		components.append(f'{d}d')
	if h != 0:
		components.append(f'{h}h')
	if m != 0:
		components.append(f'{m}m')
	if s != 0 or len(components) == 0:
		components.append(f'{s}s')
	return ''.join(components)

def describe_timedelta(td: timedelta, max_components: int = 2) -> str:
	"""
	Formats a human-readable description of a time span. E.g. "3 days 2 hours".
	"""
	d: int = td.days
	h: int = td.seconds // 3600
	m: int = (td.seconds // 60) % 60
	s: int = td.seconds % 60
	components: list[str] = []
	if d != 0:
		components.append('1 day' if d == 1 else f'{d} days')
	if h != 0:
		components.append('1 hour' if h == 1 else f'{h} hours')
	if m != 0:
		components.append('1 minute' if m == 1 else f'{m} minutes')
	if s != 0 or len(components) == 0:
		components.append('1 second' if s == 1 else f'{s} seconds')
	if len(components) > max_components:
		components = components[0:max_components]
	return ' '.join(components)

def first_command_group(cog: Cog) -> Optional[Group]:
	"""Returns the first command Group found in a cog."""
	for member_name in dir(cog):
		member = getattr(cog, member_name)
		if isinstance(member, Group):
			return member
	return None

def bot_log(guild: Optional[Guild], cog_class: Optional[Type], message: Any) -> None:
	"""Logs a message to stdout with time, cog, and guild info."""
	now: datetime = datetime.now() # local
	s = f'[{now.strftime("%Y-%m-%dT%H:%M:%S")}|'
	s += f'{cog_class.__name__}|' if cog_class else '-|'
	s += f'{guild.name}] ' if guild else '-] '
	s += str(message)
	print(s)

__QUOTE_CHARS: str = '\'"'
__ID_REGEX: re.Pattern = re.compile('^[0-9]{17,20}$')
__MENTION_REGEX: re.Pattern = re.compile('^<@[!&]([0-9]{17,20})>$')
__USER_MENTION_REGEX: re.Pattern = re.compile('^<@!([0-9]{17,20})>$')
__ROLE_MENTION_REGEX: re.Pattern = re.compile('^<@&([0-9]{17,20})>$')

def is_user_id(val: str) -> bool:
	"""Tests if a string is in user/role ID format."""
	return __ID_REGEX.match(val) is not None

def is_mention(val: str) -> bool:
	"""Tests if a string is a user or role mention."""
	return __MENTION_REGEX.match(val) is not None

def is_role_mention(val: str) -> bool:
	"""Tests if a string is a role mention."""
	return __ROLE_MENTION_REGEX.match(val) is not None

def is_user_mention(val: str) -> bool:
	"""Tests if a string is a user mention."""
	return __USER_MENTION_REGEX.match(val) is not None

def user_id_from_mention(mention: str) -> str:
	"""Extracts the user ID from a mention. Raises a ValueError if malformed."""
	m = __USER_MENTION_REGEX.match(mention)
	if m:
		return m.group(1)
	raise ValueError(f'"{mention}" is not an @ user mention')

def mention_from_user_id(user_id: Union[str, int]) -> str:
	"""Returns a Markdown user mention from a user id."""
	return f'<@!{user_id}>'

def mention_from_role_id(role_id: Union[str, int]) -> str:
	"""Returns a Markdown role mention from a role id."""
	return f'<@&{role_id}>'

def str_from_quoted_str(val: str) -> str:
	"""Removes the leading and trailing quotes from a string."""
	if len(val) < 2 or val[0:1] not in __QUOTE_CHARS or val[-1:] not in __QUOTE_CHARS:
		raise ValueError(f'Not a quoted string: {val}')
	return val[1:-1]