python-app-rocketbot/rocketbot/utils.py

"""
General utility functions.
"""
import re
import sys
import traceback
from datetime import datetime, timedelta, timezone
from typing import Any

import discord
from discord import Guild, Interaction, Permissions
from discord.app_commands import Transformer
from discord.ext.commands import BadArgument, Cog


def dump_stacktrace(e: BaseException) -> 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"

	Parameters
	----------
	s : str
		string to parse

	Returns
	-------
	timedelta

	Raises
	------
	ValueError
		if parsing fails
	"""
	p: re.Pattern = re.compile('^(?:[0-9]+[a-zA-Z])+$')
	if p.match(s) is None:
		raise ValueError(f'Illegal timespan value "{s}". Examples: 30s, 5m, 1h30m, 30d')
	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).lower()
		if unit == 'd':
			days = scalar
		elif unit == 'h':
			hours = scalar
		elif unit == 'm':
			minutes = scalar
		elif unit == 's':
			seconds = scalar
		else:
			raise ValueError(f'Invalid unit "{unit}". Valid units: "s"=seconds, "m"=minutes, "h"=hours, "d"=days')
	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 _old_first_command_group(cog: Cog) -> discord.ext.commands.Group | None:
	"""Returns the first command Group found in a cog."""
	for member_name in dir(cog):
		member = getattr(cog, member_name)
		if isinstance(member, discord.ext.commands.Group):
			return member
	return None

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

def bot_log(guild: Guild | None, cog_class: type | None, message: Any) -> None:
	"""Logs a message to stdout with time, cog, and guild info."""
	now: datetime = datetime.now(tz=None)  # noqa: DTZ005
	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})>$')
__EMAIL_REGEX: re.Pattern = re.compile(r'^(?:(?:[^<>()\[\]\\.,;:\s@"]+(?:\.[^<>()\[\]\\.,;:\s@"]+)*)|(?:".+"))@(?:(?:\[[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}])|(?:(?:[a-zA-Z\-0-9]+\.)+[a-zA-Z]{2,}))$')
__USERNAME_REGEX: re.Pattern = re.compile(r'[a-z0-9\._]{2,32}')

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 is_email_address(val: str) -> bool:
	"""Tests if a string is a well-formed email address."""
	return __EMAIL_REGEX.match(val) is not None

def is_discord_username(val: str) -> bool:
	"""Tests if a string is a properly formatted Discord username."""
	return __USERNAME_REGEX.match(val.lower())

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: str | int) -> str:
	"""Returns a Markdown user mention from a user id."""
	return f'<@!{user_id}>'

def mention_from_role_id(role_id: 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]

def blockquote_markdown(markdown: str) -> str:
	"""Encloses some Markdown in a blockquote."""
	return '> ' + (markdown.replace('\n', '\n> '))

def indent_markdown(markdown: str) -> str:
	"""Indents a block of Markdown by one level."""
	return '    ' + (markdown.replace('\n', '\n    '))

def suppress_markdown_url_previews(markdown: str) -> str:
	"""Finds URLs in markdown and encloses them in <...> to suppress the preview."""
	return re.sub(r'(?<!<)(https?://\S+)(?!>)', '<\\1>', markdown)

def truncate_markdown(markdown: str, max_length: int) -> str:
	"""Truncates markdown in a way that attempts to minimize formatting disruption."""
	if len(markdown) <= max_length:
		return markdown
	markdown = markdown[:max_length]

	# Try to cut at a newline if it's in the latter 20% of the max length
	last_newline_index = markdown.rfind('\n')
	if last_newline_index >= 0 and last_newline_index < (max_length * 8 / 10):
		return markdown[:last_newline_index] + "\n\u2026"

	# Cut at the last space
	last_space_index = markdown.rfind(' ')
	if last_space_index >= 0:
		return markdown[:last_space_index] + " \u2026"

	# Last resort, do a blind substring
	return markdown[:max_length - 1] + "\u2026"

def format_bytes(size: int) -> str:
	"""Formats s size in bytes to a human readable description (e.g. "3.2 KiB")"""
	size = max(size, 0)
	kib = 1024
	mib = kib * kib
	gib = mib * kib
	if size < kib:
		return f"{size:,} bytes"
	if size < 10 * kib:
		return f"{size/kib:,.1f} KiB"
	if size < mib:
		return f"{size/kib:,.0f} KiB"
	if size < 10 * mib:
		return f"{size/mib:,.1f} MiB"
	if size < gib:
		return f"{size/mib:,.0f} MiB"
	if size < 10 * gib:
		return f"{size/gib:,.1f} GiB"
	return f"{size/gib:,.0f} GiB"

def norm_datetime(dt: datetime) -> datetime:
	"""Converts a datetime to UTC for consistent comparison."""
	# "Naive" datetimes (without a time zone) are assumed as system local time zone
	if dt is None:
		return dt
	return datetime.fromtimestamp(dt.timestamp(), timezone.utc)

def levenshtein(a: str, b: str) -> int:
	"""Returns the Levenshtein distance between two strings."""
	# Based on https://en.wikipedia.org/wiki/Levenshtein_distance#Iterative_with_two_matrix_rows
	m: int = len(a)
	n: int = len(b)
	v0: list[int] = [ i for i in range(n + 1) ]
	v1: list[int] = [ 0 for i in range(n + 1) ]
	for i in range(m):
		v1[0] = i + 1
		for j in range(n):
			deletion_cost = v0[j + 1] + 1
			insertion_cost = v1[j] + 1
			substitution_cost = v0[j] + (0 if a[i] == b[j] else 1)
			v1[j + 1] = min(deletion_cost, insertion_cost, substitution_cost)
		h = v0
		v0 = v1
		v1 = h
	return v0[n]

MOD_PERMISSIONS: Permissions = Permissions(Permissions.manage_messages.flag)
ADMIN_PERMISSIONS: Permissions = Permissions(Permissions.administrator.flag)

class TimeDeltaTransformer(Transformer):
	async def transform(self, interaction: Interaction, value: Any) -> timedelta:
		try:
			return timedelta_from_str(str(value))
		except ValueError as e:
			print("Invalid time delta:", e)
			raise BadArgument(str(e))

	@property
	def _error_display_name(self) -> str:
		return 'timedelta'