def set_unit(unit):
    """Register a unit on a function"""

    def decorator_set_unit(func):
        func.unit = unit
        return func

    return decorator_set_unit


import math


@set_unit("cm^3")
def volume(radius, height):
    return math.pi * radius ** 2 * height


print(volume(3, 5))

print(volume.unit)
print("\n\n")


import math


def volume(radius, height) -> "cm^3":
    return math.pi * radius ** 2 * height


print(volume(3, 5))