Cleaned up code for #124