Digital root computation with benchmarking

Question

I've read this question about computing the digital root of some integer, and wondered how it would look like in Python.

As a reminder,

"If you take the digits of any number and add them together, and then add the digits of the resulting number together, and continue doing that until you get a single digit, that single digit is the digital root of the original number."

This is not much about the algorithm performance itself, since it was shown in this answer that instead of summing the digits, you could do a simple division by 9 to get the same result. The question is more about the readability, elegance, and pythonicity (if this word exists) of my code. I checked its correctness with a tool from Thonky.

More, I'd be interested in knowing your opinion on the way I deal with sys.argv, the user interaction, and the benchmarking.

Finally, is it OK to use the int type for such long numbers? It seems to work fine even with 1000-digit numbers, but intuitively it shouldn't...

Any suggestion would be highly appreciated. Thanks.

---

digitalroot.py

# -*- encoding: utf-8 -*-

import time
import random
import sys

def request_user_number():
    while True:
        try:
            nb = int(raw_input('Enter a number\n'))
            return nb
        except (ValueError, NameError):
            print 'This is not a number. Try again.'

def request_user_again():
    while True:
        try:
            ans = str(raw_input('New computation? (y/n)\n')).lower()
            if ans[0] == 'n':
                again = False
                print 'Bye !'
                break
            elif ans[0] == 'y':
                again = True
                break
            else:
                print 'Did not understand. Try again.'
        except IndexError:
            print 'Say something.'
    return again

def digits_generator(nb):
    tmp = nb
    while tmp != 0:
        yield tmp % 10
        tmp /= 10

def digital_root(nb):
    if nb < 10:
        return nb
    else:
        s = 0
        for d in digits_generator(nb):
            s += d
        return digital_root(s)

def usage():
    print 'Wrong usage'
    print '\'python digitalroot.py ui\' for a user chosen number'
    print '\'python digitalroot.py benchmark nb_size tries\' for benchmarking'
    exit(1)

if __name__ == '__main__':
    if len(sys.argv) < 2 or sys.argv[1] not in ['ui', 'benchmark']:
        usage()
    if sys.argv[1] == 'ui':
        again = True
        while again:
            nb = request_user_number()
            print 'Computing digital root of', nb
            start_time = time.clock()
            dr = digital_root(nb)
            end_time = time.clock()
            print nb, '-->', dr, '(time :', end_time-start_time, 's)'
            again = request_user_again()
    elif sys.argv[1] == 'benchmark':
        try:
            nb_size = int(sys.argv[2])
            tries = int(sys.argv[3])
        except (ValueError, NameError, IndexError):
            usage()
        global_start = time.clock()
        s = 0
        for i in xrange(tries):
            nb = random.randint(10**(nb_size-1), 10**nb_size-1)
            start_time = time.clock()
            d = digital_root(nb)
            end_time = time.clock()
            s += end_time-start_time
            if random.randint(1, tries/5) == 1:
                print 'Random control : i =', i, ',', nb, '-->', d
        print nb_size, 'digits,', tries, 'tries'
        print s, 'seconds (average =', float(s)/tries, 's/try)'
        print 'Total benchmark time :', time.clock()-global_start

Answer 1

Here are a few things:

Instead of

if ans[0] == 'n':
    again = False
    print 'Bye !'
    break
elif ans[0] == 'y':
    again = True
    break
...
return again

You could just return right away:

if ans[0] == 'n':
    print 'Bye !'
    return False
elif ans[0] == 'y':
    return True

When having a constant list, use a tuple (it makes it immediately obvious that it will not be modified, because it can't (even though it can be overridden, of course):

if len(sys.argv) < 2 or sys.argv[1] not in ['ui', 'benchmark']
if len(sys.argv) < 2 or sys.argv[1] not in ('ui', 'benchmark')

Use more functions

Currently all your code executing the things resides under if __name__ == "__main__". But since it is quite longish, I would at least put it into a main() function and only call that. But, even better, would be to put the two possible options for the problem into functions. This also allows us to define a decorator timeit to outsource the timing:

import time


def timeit(func):
    def wrapper(*arg, **kw):
        t1 = time.time()
        res = func(*arg, **kw)
        t2 = time.time()
        return (t2 - t1), res
    return wrapper


@timeit
def digital_root(nb):
    ...


def ui():
    again = True
    while again:
        nb = request_user_number()
        print 'Computing digital root of', nb
        time, dr = digital_root(nb)
        print '{} --> {} time: {}s'.format(nb, dr, time)
        again = request_user_again()


def benchmark(nb_size, tries):
    benchmark_start = time.clock()
    s = 0
    for i in xrange(tries):
        nb = random.randint(10**(nb_size-1), 10**nb_size-1)
        t, d = digital_root(nb)
        s += t
        if random.randint(1, tries/5) == 1:
            print 'Random control : i =', i, ',', nb, '-->', d
    print nb_size, 'digits,', tries, 'tries'
    print s, 'seconds (average =', float(s)/tries, 's/try)'
    print 'Total benchmark time :', time.clock()-benchmark_start


def main():
    if sys.argv[1] not in ('ui', 'benchmark'):
        usage()
    if sys.argv[1] == 'ui':
        ui()
    else:
        if len(sys.arg) != 4:
            usage()
        try:
            opts = (int(x) for x in sys.argv[2:])
        except ValueError:
            usage()
        benchmark(*opts)

if __name__ == "__main__":
    main()

We could make the main function even simpler if the two functions took the same number of arguments, allowing us to do the argument checking once and then use a dictionary for the functions. Since they are not it becomes slightly more complecated because we also need to store how many arguments we need:

def main():
    mode, opts = sys.argv[1], sys.argv[2:]
    # mode, *opts = sys.argv[1:]  # python 3.x
    func = {"ui": (ui, 0), "benchmark": (benchmark, 2)}
    if mode in func and len(opts) == func[mode][1]:
        try:
            opts = (int(opt) for opt in opts)
        except ValueError:
            usage()
        func[mode][0](*opts)
    else:
        usage()

Alternatively, it is easier to ask forgiveness than permission:

def main():
    mode, opts = sys.argv[1], sys.argv[2:]
    # mode, *opts = sys.argv[1:]  # python 3.x
    func = {"ui": ui, "benchmark": benchmark}
    try:
        opts = (int(opt) for opt in opts)
        func[mode](*opts)
    except (KeyError, TypeError, ValueError):  # KeyError for wrong func name, TypeError for nopts mismatch, ValueError if opts not numbers
            usage()

Where it is arguably whether the except might not match too many different failure scenarios, so it might be better to go with one of the more explicit approaches above.

Answer 2

I think your algorithmic part is exceedingly verbose. One of the nicenesses of Python is how it can lead to code that reads almost like English. And for this particular problem, even leaving performance considerations aside, I think iteration is clearer than recursion. So I would rewrite things as:

def digits_of(number):
    """Yields the digits of an integer."""
    while number > 0:
        yield number % 10
        number //= 10

def digital_root(number):
    """Computes the digital root of an integer."""
    while number > 9:
        number = sum(digits_of(number))
    return number

I'm pretty sure my 6 year old daughter could make sense of that last function with very little help.

Answer 3

Regarding digits_generator

You don't need to use a new variable. You could use nb directly.

You may use the fact that non-0 integers are considered as True in a boolean context and 0 is considered as False to write if nb:.

Anytime you use % and / on the same values, you could use divmod. Here you could write:

def digits_generator2(nb):
    while nb:
        q, r = divmod(nb, 10)
        yield r
        nb = q

or:

def digits_generator2(nb):
    while nb:
        nb, r = divmod(nb, 10)
        yield r

You could avoid having the magic number 10 in your code by using a parameter with a meaningful name and a useful default value like:

def digits_generator2(nb, base=10):
    while nb:
        nb, r = divmod(nb, base)
        yield r