Using doctest module for if __name__ == '__main__' test

Many people including me write a test code in __name__ == '__main__' like

def somefunction(intvalue):
 return intvalue * 2

if __name__ == '__main__':
 #Tests whatever.
 if somefunction(2) != 4:
     print "NG"

Previsously I did something like this.

if __name__ == '__main__':

 def test(cmd):
     print cmd,
     exec cmd

 ph = PreferenceHolder(pref1=1.0, pref2=100, pref3='tamtam')
 prefs = ph.getPreferences().items()
 test("print prefs == [('pref1', 1.0), ('pref2', 100), ('pref3', 'tamtam')]")
 test("print ph.pref1 == 1.0")

I wanted a nicer way so looked around documents for Python unit test modules. Python has two modules for unit testing, unittest and doctest. unittest is kind of overkill for a test in __name__ == '__main__' so I looked at doctest, which is better because it's purpose is to offer tools to write "readable" tests. I serched for a function to let me write a test code like

if __name__ == '__main__':
 teststr = """
 >>> ph = PreferenceHolder(pref1=1.0, pref2=100, pref3='tamtam')
 >>> ph.getPreferences().items()
 [('pref1', 1.0), ('pref2', 100), ('pref3', 'tamtam')]
 >>> ph.pref1
 1.0
 """
 import doctest
 doctest.teststring(teststr) #Note: it doesn't exist

Unfortunately I found I cannot do this. I could find a function that takes a filename to get a docstring but couldn't find one that takes a Python string directly like above code. I think this function is not provided because doctest module is designed to write the whole test in docstrings (It is provided as a class but not as a userfriendly function. Readability is important because everybody looks at this test code as an example). I was just looking for an easy way to give the users usage examples (primary interest) and have it work as a test (secondary), so I didn't want to write the whole tests in each doc string. I found an alternative function run_docstring_examples() which tests only a docstring for a specified object given by one of its arguments. The document says

There’s also a function to run the doctests associated with a single object. This function is provided for backward compatibility. There are no plans to deprecate it, but it’s rarely useful

but it's nearly perfect for my purpose. So now, I can write the test like this.

if __name__ == '__main__':
 def test():
     """
     >>> ph = PreferenceHolder(pref1=1.0, pref2=100, pref3='tamtam')
     >>> ph.getPreferences().items()
     [('pref1', 1.0), ('pref2', 100), ('pref3', 'tamtam')]
     >>> ph.pref1
     1.0
     """
 import doctest
 doctest.run_docstring_examples(test, globals(), False, __name__)

Because it's a doctest, it prints nothing if all the tests finish successfully (unless you set True to the third argument of run_docstring_examples). If it fails, it prints an error message

$ python preferenceholder.py
**********************************************************************
File "preferenceholder.py", line 94, in __main__
Failed example:
 ph.pref1
Expected:
 0.0
Got:
 1.0
$

Nice isn't it?
Please see the previous blog entry for the actual usage.

By the way, one common pitfall of using doctest is you don't update comments in a docstring when you modify its implementation. I've seen lots of wrong comments in existing dectests. Try to write only comments about specs and not about implementation details unless it is necessary.