Actually added the code mentioned in my last commit. Whoops.