Задача с использованием функции-генератора

user_input = int(input("Сколько аргументов использовать? --> "))
if user_input == 1:
    a = int(input("от "))
    print(*(i for i in range(a)))
elif user_input == 2:
    a,b = int(input("от ")), int(input("до: "))
    print(*(i for i in range(a, b +1)))
elif user_input == 3:
    a,b,c = int(input("от: ")), int(input("до: ")),int(input("с шагом: "))
    print(*(i for i in range(a, b +1, c)))
 
else:
    print("Неверный ввод")

def ge(start, stop, step):
    if step < 0:
        i = start
        while i > stop:
            yield i
            i += step
    elif step > 0:
        i = start
        while i < stop:
            yield i
            i += step
 
 
def range_generator(*args):
    match args:
        case (stop, ):
            start, step = 0, 1
        case (start, stop):
            step = 1
        case (start, stop, step) if step != 0:
            pass
        case _:
            raise ValueError('step must not be zero')
    return ge(start, stop, step)
 
 
try:
    print(*range_generator(5))
    print(*range_generator(5, 23))
    print(*range_generator(5, 23, 2))
    print(*range_generator(12, 1, -2))
    print(*range_generator(12, 10))
    print(*range_generator(20, 10, 0))
except ValueError as ve:
    print(ve)

def range_generator(*args):
    amount_args = len(args)
    if amount_args == 1:
        start, step = 0, 1
        stop, *_ = args
    elif amount_args == 2:
        start, stop, *_ = args
        step = 1
    else:
        start, stop, step, *_ = args
        if step == 0:
            raise ValueError('step must not be zero')
    return ge(start, stop, step)