Will submit a new PR for this.