from typing import Iterable, Optional

from machine.annotations.range import Range
from machine.tokenization import StringTokenizer

from machine.tokenization import WhitespaceTokenizer

class WhitespaceIncludedTokenizer(StringTokenizer):
    def tokenize_as_ranges(self, data: str, data_range: Optional[Range[int]] = None) -> Iterable[Range[int]]:
        if data_range is None:
            data_range = Range.create(0, len(data))
        start_index = -1
        for i in data_range:
            if self._is_whitespace(data[i]):
                if start_index != -1:
                    yield Range.create(start_index, i)
                yield Range.create(i, i+1) # Added for kathairo.py to include whitespace
                start_index = -1
            elif start_index == -1:
                start_index = i

        if start_index != -1:
            yield Range.create(start_index, data_range.end)

    def _is_whitespace(self, c: str) -> bool:
        return c.isspace() or c == "\u200b" or c == "\ufeff"


WHITESPACE_TOKENIZER = WhitespaceTokenizer()
