Passing a color argument to keys() for custom highlight color

Hello there! I have a custom highlight class that allows highlighting with different colours. I managed to do it taking an idea from one of the posts here: Create multiple instances of the class with different colours and dynamically change them, passing a color on command call.

Having done that, I’ve faced an issue: keys() uses the first highlight class in a row (I have an array of those highlight classes) instead of calling the one with the appropriate color.

So my question is: Is there a way to pass the color as I did with command.highlight(color) ? Or, maybe, if’ve got it totally wrong and there’s a much easier way to implement it.

Thanks in advance <3