Cleaning up code and added better attribute handling