Fix imports errors caused by previous commit.

Signed-off-by: Stephen Finucane <stephenfinucane@hotmail.com>