A simple script to compute the Collatz sequence (aka the 3n+1 problem), bending the rules a bit to let real numbers into the game

collatz.py

#!/usr/bin/env python3
## This looks like a fun problem
import argparse
from typing import Optional, Union


def gen(n: Union[int, float]) -> Union[int, float]:
    """
    Handy generator for the 3n+1
    sequence
    """

    while True:
        if n % 2 == 0:
            n = n // 2
            yield n
        else:
            n  = n * 3 + 1
            yield n


def main(start: Union[int, float], limit: Optional[int] = None,
         verbose: bool = False) -> int:
    """
    Computes the 3n+1 sequence for
    a given number (can be float or int,
    positive or negative) and returns its
    stepping distance. If the stepping
    distance exceeds the given limit,
    which defaults to None (no limit) the
    sequence is interrupted
    """

    for i, n in enumerate(gen(start), 1):
        if verbose:
            print(f"Iteration {i}: {n}", flush=True)
        if limit and i == limit:
            break
        elif n > 0 and n == 1:
            # We reached the end on the
            # positive graph
            break
        elif n < 0 and n in {-17, -5, -2}:
            # End of the negative graph
            break
    return i


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("n", help="The number for which the 3n+1 sequence" \
                                  " should be calculated, can be int or float")
    parser.add_argument("--limit", help="The limit to the stepping distance" \
                                        " of the sequence. Defaults to no limit",
                        type=int, default=None)
    parser.add_argument("--verbose", help="Prints each number in the sequence",
                        action="store_true")
    parser.add_argument("--cast-to", help="Force casting of n to the given type" \
                                          ", can be 'int' or 'float'. By default the type is inferred automatically",
                        choices=("int", "float"), default=None)
    args = parser.parse_args()
    n = args.n.lower()
    if "." not in n and "e" not in n:
        try:
            args.n = int(n)
        except ValueError:
            print("Error: Invalid number literal for n")
            exit()
    else:
        try:
            args.n = float(n)
        except ValueError:
            print("Error: Invalid number literal for n")
            exit()
    if args.cast_to:
        args.n = int(n) if args.cast_to == "int" else float(n)
    print(f"The stepping distance of the 3n+1 sequence for {args.n} (limited to {args.limit}) is", main(args.n, args.limit, args.verbose))