How To Write Doctests in Python

Introduction

Documentation and testing are core components of every productive software development process. Ensuring that code is thoroughly documented and tested not only ensures that a program runs as expected, but also supports collaboration across programmers as well as user adoption. A programmer can be well served by first writing documentation and then tests, before finally writing code. Following a process like this will ensure that the function one is coding (for example) is well thought out and addresses possible edge cases.

Python’s standard library comes equipped with a test framework module called doctest. The doctest module programmatically searches Python code for pieces of text within comments that look like interactive Python sessions. Then, the module executes those sessions to confirm that the code referenced by a doctest runs as expected.

Additionally, doctest generates documentation for our code, providing input-output examples. Depending on how you approach writing doctests, this can either be closer to either “‘literate testing’ or ‘executable documentation,’” as the Python Standard Library documentation explains.

Prerequisites

You should have Python 3 installed and a programming environment set up on your computer or server. If you don’t have a programming environment set up, you can refer to the installation and setup guides for a local programming environment or for a programming environment on your server appropriate for your operating system (Ubuntu, CentOS, Debian, etc.)

Doctest Structure

A Python doctest is written as though it is a comment, with a series of three quotation marks in a row — """ — at the top and bottom of the doctest.

Sometimes, doctests are written with an example of the function and the expected output, but it may be preferable to also include a comment on what the function is intended to do. Including a comment will ensure that you as the programmer have sharpened your goals, and a future person reading the code understands it well. Remember, the future programmer reading the code may very well be you.

The following is a mathematical example of a doctest for a function such as add(a, b) that adds two numbers together:

"""
Given two integers, return the sum.

>>> add(2, 3)
5
"""

In this example we have a line of explanation, and one example of the add() function with two integers for input values. If in the future you want the function to be able to add more than two integers, you would need to revise the doctest to match the function’s inputs.

So far, this doctest is very readable to a human. You can further iterate on this docstring by including machine-readable parameters and a return description to explicate each variable coming in and out of the function.

Here, we’ll add docstrings for the two arguments that are passed to the function and the value that is returned. The docstring will note the data types for each of the values — the parameter a, the parameter b, and the returned value — in this case they are all integers.

"""
Given two integers, return the sum.

:param a: int
:param b: int
:return: int

>>> add(2, 3)
5
"""

This doctest is now ready to be incorporated into a function and tested.

Incorporating a Doctest into a Function

Doctests sit within a function after the def statement and before the code of the function. As this follows the initial definition of the function, it will be indented following Python’s conventions.

This short function indicates how the doctest is incorporated.

def add(a, b):
    """
    Given two integers, return the sum.

    :param a: int
    :param b: int
    :return: int

    >>> add(2, 3)
    5
    """
    return a + b

In our short example, we only have this one function in our program, so now we will have to import the doctest module and have a call statement for the doctest to run.

We’ll be adding the following lines before and after our function:

import doctest 
...
doctest.testmod()

At this point, let’s test it on the Python shell rather than saving it to a program file right now. You can access a Python 3 shell on your command line terminal of choice (including IDE terminal) with the python3 command (or python if you’re using a virtual shell).

python3

If you go this route, once you press ENTER, you’ll receive output similar to the following:

Output
Type "help", "copyright", "credits" or "license" for more information.
>>> 

You’ll be able to start typing code after the >>> prompt.

Our complete example code, including the add() function with a doctest, docstrings, and a call to invoke the doctest follows. You can paste it into your Python interpreter to try it out:

import doctest


def add(a, b):
    """
    Given two integers, return the sum.

    :param a: int
    :param b: int
    :return: int

    >>> add(2, 3)
    5
    """
    return a + b

doctest.testmod()

Once you run the code, you’ll receive the following output:

Output
TestResults(failed=0, attempted=1)

This means that our program ran as expected!

If you modify the program above so that the return a + b line is instead return a * b, which modifies the function to multiply integers and return their product instead, you’ll receive a failure notice:

Output
**********************************************************************
File "__main__", line 9, in __main__.add
Failed example:
    add(2, 3)
Expected:
    5
Got:
    6
**********************************************************************
1 items had failures:
   1 of   1 in __main__.add
***Test Failed*** 1 failures.
TestResults(failed=1, attempted=1)

From the output above, you can begin to understand how useful the doctest module is as it fully describes what happened when a and b were multiplied instead of added, returning the product of 6 in the example case.

You may want to experiment with more than one example. Let’s try with an example where both variables a and b contain the value of 0, and change the program back to addition with the + operator.

import doctest


def add(a, b):
    """
    Given two integers, return the sum.

    :param a: int
    :param b: int
    :return: int

    >>> add(2, 3)
    5
    >>> add(0, 0)
    0
    """
    return a + b

doctest.testmod()

Once we run this, we’ll receive the following feedback from the Python interpreter:

Output
TestResults(failed=0, attempted=2)

Here, the output indicates that the doctest attempted two tests, on both lines of add(2, 3) and add(0, 0) and that both passed.

If, again, we change the program to use the * operator for multiplication rather than the + operator, we can learn that edge cases are important when working with the doctest module, because the second example of add(0, 0) will return the same value whether it is addition or multiplication.

import doctest


def add(a, b):
    """
    Given two integers, return the sum.

    :param a: int
    :param b: int
    :return: int

    >>> add(2, 3)
    5
    >>> add(0, 0)
    0
    """
    return a * b

doctest.testmod()

The following output is returned:

Output
**********************************************************************
File "__main__", line 9, in __main__.add
Failed example:
    add(2, 3)
Expected:
    5
Got:
    6
**********************************************************************
1 items had failures:
   1 of   2 in __main__.add
***Test Failed*** 1 failures.
TestResults(failed=1, attempted=2)

When we modify the program, only one of the examples fails, but it is fully described as before. If we had started with the add(0, 0) example rather than the add(2, 3) example, we may not have noticed that there were opportunities for failure when small components of our program change.

Doctests in Programming Files

So far, we have used an example on the Python interactive terminal. Let’s now use it in a programming file that will count the number of vowels in a single word.

In a program, we can import and call the doctest module in our if __name__ == "__main__": clause at the bottom of our programming file.

We’ll create a new file — counting_vowels.py — in our text editor, you can use nano on the command line, like so:

nano counting_vowels.py

We can begin with defining our function count_vowels and passing the parameter of word to the function.

def count_vowels(word):

Before we write the body of the function, let’s explain what we want the function to do in our doctest.

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

So far so good, we are being pretty specific. Let’s flesh this out with the data type of the parameter word and the data type we want returned. In the first case it’s a string, in the second case it’s an integer.

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

    :param word: str
    :return: int

Next, let’s find examples. Think of a single word that has vowels, and then type it into the docstring.

Let’s choose the word 'Cusco' for the city in Peru. How many vowels are in “Cusco”? In English, vowels are often considered to be a, e, i, o, and u. So here we will count u and o as the vowels.

We’ll add the test for Cusco and the return of 2 as the integer into our program.

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

    :param word: str
    :return: int

    >>> count_vowels('Cusco')
    2

Again, it’s a good idea to have more than one example. Let’s have another example with more vowels. We’ll go with 'Manila' for the city in the Philippines.

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

    :param word: str
    :return: int

    >>> count_vowels('Cusco')
    2

    >>> count_vowels('Manila')
    3
    """

Those doctests look great, now we can code our program.

We’ll start with initializing a variable — total_vowels to hold the vowel count. Next, we’ll create a for loop to iterate across the letters of the word string, and then include a conditional statement to check whether each letter is a vowel. We’ll increase the vowel count through the loop, then return the total number of vowels in the word to the total_values variable. Our program should be similar to this, without the doctest:

def count_vowels(word):
    total_vowels = 0
    for letter in word:
        if letter in 'aeiou':
            total_vowels += 1
    return total_vowels

If you need more guidance on these topics, please check out our How To Code in Python book or complementary series.

Next, we’ll add our main clause at the bottom of the program and import and run the doctest module:

if __name__ == "__main__":
    import doctest
    doctest.testmod()

At this point, here is our program:

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

    :param word: str
    :return: int

    >>> count_vowels('Cusco')
    2

    >>> count_vowels('Manila')
    3
    """
    total_vowels = 0
    for letter in word:
        if letter in 'aeiou':
            total_vowels += 1
    return total_vowels

if __name__ == "__main__":
    import doctest
    doctest.testmod()

We can run the program by using the python (or python3 depending on your virtual environment) command:

python counting_vowels.py

If your program is identical to the above, all the tests should have passed and you will not receive any output. This means that the tests passed. This silent feature is useful when you are running programs for other purposes. If you are running specifically to test, you may want to use the -v flag, as in the following:

python counting_vowels.py -v

When you do, you should receive this output:

Output
Trying:
    count_vowels('Cusco')
Expecting:
    2
ok
Trying:
    count_vowels('Manila')
Expecting:
    3
ok
1 items had no tests:
    __main__
1 items passed all tests:
   2 tests in __main__.count_vowels
2 tests in 2 items.
2 passed and 0 failed.
Test passed.

Excellent! The test has passed. Still, our code may not be quite optimized for all edge cases yet. Let’s learn how to use doctests to strengthen our code.

Using Doctests to Improve Code

At this point, we have a working program. Maybe it is not the best program it can be yet, so let’s try to find an edge case. What if we add an upper-case vowel?

Add another example in the doctest, this time let’s try 'Istanbul' for the city in Turkey. Like Manila, Istanbul also has three vowels.

Here is our updated program with the new example:

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

    :param word: str
    :return: int

    >>> count_vowels('Cusco')
    2

    >>> count_vowels('Manila')
    3

    >>> count_vowels('Istanbul')
    3
    """
    total_vowels = 0
    for letter in word:
        if letter in 'aeiou':
            total_vowels += 1
    return total_vowels

if __name__ == "__main__":
    import doctest
    doctest.testmod()

Let’s run the program again.

python counting_vowels.py

We have identified an edge case! This is the output we have received:

Output
**********************************************************************
File "counting_vowels.py", line 14, in __main__.count_vowels
Failed example:
    count_vowels('Istanbul')
Expected:
    3
Got:
    2
**********************************************************************
1 items had failures:
   1 of   3 in __main__.count_vowels
***Test Failed*** 1 failures.

The output above indicates that the test on 'Istanbul' is the one that failed. We told the program we were expecting three vowels to be counted, but instead the program counted only two. What went wrong here?

In our line if letter in 'aeiou': we have only passed in lower-case vowels. We can modify our 'aeiou' string to be 'AEIOUaeiou' to count both upper- and lower-case vowels, or we can do something more elegant, and convert our value stored in word to lower-case with word.lower(). Let’s do the latter.

def count_vowels(word):
    """
    Given a single word, return the total number of vowels in that single word.

    :param word: str
    :return: int

    >>> count_vowels('Cusco')
    2

    >>> count_vowels('Manila')
    3

    >>> count_vowels('Istanbul')
    3
    """
    total_vowels = 0
    for letter in word.lower():
        if letter in 'aeiou':
            total_vowels += 1
    return total_vowels

if __name__ == "__main__":
    import doctest
    doctest.testmod()

Now, when we run the program, all tests should pass. You can confirm again by running python counting_vowels.py -v with the verbose flag.

Still, this probably is not the best program it can be, and it may not be considering all edge cases.

What if we pass the string value 'Sydney' — for the city in Australia — to word? Would we expect three vowels or one? In English, y is sometimes considered to be a vowel. Additionally, what would happen if you use the value 'Würzburg' — for the city in Germany — would the 'ü' count? Should it? How will you handle other non-English words? How will you handle words that use different character encodings, such as those available in UTF-16 or UTF-32?

As a software developer, you will sometimes need to make tricky decisions like deciding which characters should be counted as vowels in the example program. Sometimes there may not be a right or wrong answer. In many cases, you will not consider the full scope of possibilities. The doctest module is therefore a good tool to start to think through possible edge cases and capture preliminary documentation, but ultimately you will need human user testing — and very likely collaborators — to build robust programs that serve everyone.

Conclusion

This tutorial introduced the doctest module as not only a method for testing and documenting software, but also as a way to think through programming before you begin, by first documenting it, then testing it, then writing the code.

Not writing tests could lead not only to bugs but software failure. Getting in the habit of writing tests prior to writing code can support productive software that serves other developers and end users alike.

If you would like to learn more about testing and debugging, check out our “Debugging Python Programs” series. We also have a free eBook on How To Code in Python and another on Python Machine Learning Projects.