Minor cleanup to improve pythonicity