Fix issues where it was splitting def / class definitions