Your proposed solution sounds great. It reuses your existing work and only adds a small amount of complexity to the logic. There might be a more efficient way, but I won't take the time to look for one.