Fix missing import.

This is cherry-picked from the commit in master that addressed
a bunch of pyflakes warnings, however this import missing
is not innocuous.  It really does break stuff.