I've thought about it and looked through the various functions/classes. I think @alexgibson's idea will work and make it more robust in the future. I'll try to make those changes today or tomorrow.