Revert "Slight edit to py code"

This reverts commit 753e954e883f89da54008f8f75153fe930da0405.