Feature/change add class check

Hello again. 

I made the changes you suggested tested on a build of mine. Ran gulp. Hopefully all good to go. 

Cheers.