python-app-rocketbot/rocketbot/utils.py

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

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

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]