Fix python extension build

Oops, forgot to fix the python extension