Partitioning integer `n` using parts up to size `m`

integer_partitions.py

#!/usr/bin/env python3

"""
Problem statement:
- https://composingprograms.com/pages/17-recursive-functions.html#example-partitions
- https://youtu.be/ngCos392W4w?t=678
"""


from functools import lru_cache
from typing import Generator


@lru_cache
def count_partitions(n: int, m: int) -> int:
    """
    Recursive implementation of partition counter
    """
    # Base case(s)
    if n == 0:
        return 1
    if m == 0 or n < 0:
        return 0

    # Recursive case
    return count_partitions(n, m - 1) + count_partitions(n - m, m)


def integer_partitions(n: int, m: int) -> Generator:
    """
    Print ways to partition integer `n` using integers up to and including `m`
    using generator function
    """
    if n < 0 or m == 0:
        return
    if n == m:
        yield str(m)
    for p in integer_partitions(n - m, m):
        yield f'{p} + {m}'
    yield from integer_partitions(n, m - 1)


def main():
    """
    Entrypoint for partition counter program
    """
    while True:
        # Handle user I/O
        is_error = True
        while is_error:
            try:
                n = int(input("n: "))
                m = int(input("m: "))
                if n < 0 or m < 0:
                    raise ValueError
                is_error = False
            except Exception:
                print("User input cannot be casted to integer! Try again.\n")

        # Run program for valid inputs
        answer = count_partitions(n, m)
        print(f"{answer} count_partitions\n")
        for p in integer_partitions(n, m):
            print(f"{p}", f"= {n}")
        print()


if __name__ == "__main__":
    main()