Make connection code more pythonic