Redesign ensure_annotations wrapper