Fix import statement in python driver

Closes #1438
Review over IRC by @neumino