Writing N-body gravity simulations code in Python