Add extension method instead of inline TODO