Change implementation to Python