Merge pull request #2 from alistairhill/fix_logic

Fix logic