Minor changes

Hey hey!

I updated your script since I like it and it wasn't doing exacly what I wanted :smile: 

I hope these changes make sense.