DnD ability scores roller helper

dnd_scores.py

"""DnD ability scores roller helper.

CC0 - Credit appreciated.
by HacKan 2024 - https://hackan.net

Requirements:
    - Python 3.10+
"""

import argparse
from collections.abc import Iterable
from secrets import randbelow
from typing import Final, TypeAlias

AbilityScoreT: TypeAlias = int
AbilityScoresT: TypeAlias = tuple[
    AbilityScoreT,
    AbilityScoreT,
    AbilityScoreT,
    AbilityScoreT,
    AbilityScoreT,
    AbilityScoreT,
]
RollT: TypeAlias = int
RollsT: TypeAlias = tuple[RollT, RollT, RollT, RollT]

STANDARD_ARRAY: Final[AbilityScoresT] = (15, 14, 13, 12, 10, 8)
DICE_CHARS: Final = ('⚀', '⚁', '⚂', '⚃', '⚄', '⚅')
QUALIFICATIONS: Final = ('👎', '🤏', '👌', '👍', '🤘', '🚀')


def roll_d6() -> int:
    """Roll a d6 die."""
    return randbelow(6) + 1


def roll_ability() -> RollsT:
    """Roll 4d6 for an ability score."""
    return tuple(sorted((roll_d6() for _ in range(4)), reverse=True))


def ability_score_from_roll(roll: RollsT, /) -> AbilityScoreT:
    return sum(roll[:3])


def ability_qualification(ability: AbilityScoreT, /) -> str:
    if ability > 16:
        return QUALIFICATIONS[-1]

    if ability > 14:
        return QUALIFICATIONS[-2]

    if ability > 12:
        return QUALIFICATIONS[-3]

    if ability > 9:
        return QUALIFICATIONS[-4]

    if ability > 7:
        return QUALIFICATIONS[-5]

    return QUALIFICATIONS[0]


def overall_abilities_qualification(abilities: AbilityScoresT) -> str:
    sums = sum(abilities)
    std_sums = sum(STANDARD_ARRAY)

    if sums > std_sums * 1.2:
        return QUALIFICATIONS[-1]

    if sums > std_sums * 1.1:
        return QUALIFICATIONS[-2]

    if sums > std_sums:
        return QUALIFICATIONS[-3]

    if sums > std_sums * 0.9:
        return QUALIFICATIONS[-4]

    if sums > std_sums * 0.8:
        return QUALIFICATIONS[-5]

    return QUALIFICATIONS[0]


def analyze_score_or_rolls(
    *,
    num: int,
    score: AbilityScoreT,
    roll: RollsT | None,
) -> None:
    """Print a summary of the analysis of given score and roll, if any."""
    if roll:
        roll_text = f': {" ".join(DICE_CHARS[die - 1] for die in roll)}'
    else:
        roll_text = ''

    assert score

    print(
        f'Roll #{num}{roll_text}',
        f'=> Ability score: {score: <2}',
        ability_qualification(score),
    )


def analyze_scores(scores: Iterable[AbilityScoreT], /) -> None:
    """Analyze given scores and print a summary."""
    print('Analyzing given ability scores...')
    print('Qualifiers:', ', '.join(QUALIFICATIONS))
    print()
    for num, score in enumerate(scores, start=1):
        analyze_score_or_rolls(num=num, score=score, roll=None)


def roll_ability_scores() -> AbilityScoresT:
    """Roll ability scores and print a summary.

    The method used is rolling 4d6 and using the highest 3 for each ability.
    """
    print('Rolling random ability scores...')
    print('Rolling 4 dice and choosing the highest 3...')
    print('Qualifiers:', ', '.join(QUALIFICATIONS))
    print()
    scores = []
    for num in range(1, 7):
        roll = roll_ability()
        score = ability_score_from_roll(roll)
        scores.append(score)
        analyze_score_or_rolls(num=num, score=score, roll=roll)

    scores.sort(reverse=True)

    return tuple(scores)


def parse_args() -> None:
    """Parse CLI arguments."""
    parser = argparse.ArgumentParser(
        prog='dnd_scores',
        description=(
            'DnD 5E Ability Scores roll or analyze.\n'
            + 'Input ability scores to analyze, or run the script to generate random '
            + 'ability scores using the method 3 out of 4d6'
        ),
    )

    parser.add_argument('scores', type=int, nargs='*', help='Scores to analyze')

    args = parser.parse_args()

    return args


def main() -> None:
    print('DnD Ability Scores Helper')
    print()

    args = parse_args()

    if scores := args.scores:
        if len(scores) < 6:
            warn = 'few'
        elif len(scores) > 6:
            warn = 'many'
        else:
            warn = ''

        if warn:
            print('Warning: too', warn, 'ability scores given (should be 6)!')

        analyze_scores(scores)
    else:
        scores = roll_ability_scores()

    print()
    print(
        'Ability scores:',
        ', '.join(str(score) for score in sorted(scores, reverse=True)),
        overall_abilities_qualification(scores),
    )
    print()
    print('Have fun!')


if __name__ == '__main__':
    main()