Minor fixes to Python wrapper