Fix missing imports.

These were found with `pylint -E`.
Merge pull request #3 from mgrundy/kamran/fix-missing-imports

Fix missing imports.