Closing #160, @mikeshulman should check if these changes are ok.