<commit_before>import typing as T
from collections import defaultdict
from collections.abc import Mapping

from conllu.exceptions import ParseException
from conllu.serializer import serialize

DEFAULT_EXCLUDE_FIELDS = ('id', 'deprel', 'xpos', 'feats', 'head', 'deps', 'misc')

if T.TYPE_CHECKING:
    class SupportsIndex(T.Protocol):
        def __index__(self) -> int:
            pass

    class SupportsNext(T.Protocol):
        def __next__(self) -> int:
            pass

class Metadata(dict):
    pass

class Token(dict):
    MAPPING = {
        "upostag": "upos",
        "xpostag": "xpos",
        "upos": "upostag",
        "xpos": "xpostag",
    }

    def get(self, key: str, default: T.Optional[T.Any] = None) -> T.Any:
        if key not in self and key in self.MAPPING:
            key = self.MAPPING[key]

        return super(Token, self).get(key, default)

    def __missing__(self, key: str) -> T.Any:
        if key in self.MAPPING:
            return self.get(self.MAPPING[key])

        raise KeyError("'" + key + "'")

    def __str__(self) -> str:
        if 'form' in self:
            return self['form']

        if 'id' in self:
            return f"id={self['id']}"

        return ''

class TokenList(T.List[Token]):
    def __init__(
        self,
        tokens: T.Iterable[Token] = None,
        metadata: Metadata = None,
        default_fields: T.Optional[T.Iterable[str]] = None,
    ):
        tokens = tokens or []

        if not isinstance(tokens, list):
            raise ParseException("Can't create TokenList, tokens is not a list.")

        if len(tokens) > 0 and not isinstance(tokens[0], Token):
            tokens = [Token(token) for token in tokens]

        super(TokenList, self).__init__(tokens)

        self.metadata = metadata or Metadata()
        self.default_fields = default_fields

    def __repr__(self) -> str:
        tokens = ', '.join(str(token) for token in self)
        if not self.metadata:
            return f'TokenList<{tokens}>'
        else:
            metadata = ', '.join(f"{key}: \"{value}\"" for key, value in self.metadata.items())
            return f'TokenList<{tokens}, metadata={{{metadata}}}>'

    def __eq__(self, other: T.Any) -> bool:
        if not isinstance(other, TokenList):
            other = TokenList(other)

        return super(TokenList, self).__eq__(other) and self.metadata == other.metadata

    def __ne__(self, other: T.Any) -> bool:
        return not self == other

    def clear(self) -> None:
        super(TokenList, self).clear()
        self.metadata = Metadata()
        self.default_fields = None

    def copy(self) -> 'TokenList':
        tokens_copy = super().copy()
        return TokenList(tokens_copy, self.metadata, self.default_fields)

    def extend(self, iterable: T.Union['TokenList', T.Iterable[Token]]) -> None:
        if not isinstance(iterable, TokenList):
            iterable = TokenList(iterable)

        super(TokenList, self).extend(iterable)

        self.metadata.update(iterable.metadata)

    def _dict_to_token_and_set_defaults(self, token: T.Union[dict, Token]) -> Token:
        if not isinstance(token, Token):
            token = Token(token)

        if self.default_fields:
            for field in self.default_fields:
                if field not in token:
                    token[field] = "_"

        return token

    def append(self, token: T.Union[dict, Token]) -> None:
        token = self._dict_to_token_and_set_defaults(token)
        super(TokenList, self).append(token)

    def insert(self, i: 'SupportsIndex', token: T.Union[dict, Token]) -> None:
        token = self._dict_to_token_and_set_defaults(token)
        super(TokenList, self).insert(i, token)

    @T.overload
    def __setitem__(self, key: 'SupportsIndex', tokens: T.Union[dict, Token]) -> None: ...  # noqa, pragma: no cover

    @T.overload
    def __setitem__(self, key: slice, tokens: T.Union[T.Iterable[T.Union[dict, Token]], 'TokenList']) -> None: ...  # noqa, pragma: no cover

    def __setitem__(self, key, tokens):  # noqa: F811
        if isinstance(key, slice):
            tokens = [self._dict_to_token_and_set_defaults(token) for token in tokens]
            super(TokenList, self).__setitem__(key, tokens)
        else:
            token = tokens
            token = self._dict_to_token_and_set_defaults(token)<commit_msg>Carefully adjust type definitions to keep mypy happy.<commit_after>import typing as T
from collections import defaultdict
from collections.abc import Mapping

from conllu.exceptions import ParseException
from conllu.serializer import serialize

DEFAULT_EXCLUDE_FIELDS = ('id', 'deprel', 'xpos', 'feats', 'head', 'deps', 'misc')

if T.TYPE_CHECKING:
    class SupportsIndex(T.Protocol):
        def __index__(self) -> int:
            pass

    class SupportsNext(T.Protocol):
        def __next__(self) -> int:
            pass

class Metadata(dict):
    pass

class Token(dict):
    MAPPING = {
        "upostag": "upos",
        "xpostag": "xpos",
        "upos": "upostag",
        "xpos": "xpostag",
    }

    def get(self, key: str, default: T.Optional[T.Any] = None) -> T.Any:
        if key not in self and key in self.MAPPING:
            key = self.MAPPING[key]

        return super(Token, self).get(key, default)

    def __missing__(self, key: str) -> T.Any:
        if key in self.MAPPING:
            return self.get(self.MAPPING[key])

        raise KeyError("'" + key + "'")

    def __str__(self) -> str:
        if 'form' in self:
            return self['form']

        if 'id' in self:
            return f"id={self['id']}"

        return ''

class TokenList(T.List[Token]):
    def __init__(
        self,
        tokens: T.Iterable[Token] = None,
        metadata: Metadata = None,
        default_fields: T.Optional[T.Iterable[str]] = None,
    ):
        tokens = tokens or []

        if not isinstance(tokens, list):
            raise ParseException("Can't create TokenList, tokens is not a list.")

        if len(tokens) > 0 and not isinstance(tokens[0], Token):
            tokens = [Token(token) for token in tokens]

        super(TokenList, self).__init__(tokens)

        self.metadata = metadata or Metadata()
        self.default_fields = default_fields

    def __repr__(self) -> str:
        tokens = ', '.join(str(token) for token in self)
        if not self.metadata:
            return f'TokenList<{tokens}>'
        else:
            metadata = ', '.join(f"{key}: \"{value}\"" for key, value in self.metadata.items())
            return f'TokenList<{tokens}, metadata={{{metadata}}}>'

    def __eq__(self, other: T.Any) -> bool:
        if not isinstance(other, TokenList):
            other = TokenList(other)

        return super(TokenList, self).__eq__(other) and self.metadata == other.metadata

    def __ne__(self, other: T.Any) -> bool:
        return not self == other

    def clear(self) -> None:
        super(TokenList, self).clear()
        self.metadata = Metadata()
        self.default_fields = None

    def copy(self) -> 'TokenList':
        tokens_copy = super().copy()
        return TokenList(tokens_copy, self.metadata, self.default_fields)

    def extend(self, iterable: T.Union['TokenList', T.Iterable[Token]]) -> None:
        if not isinstance(iterable, TokenList):
            iterable = TokenList(iterable)

        super(TokenList, self).extend(iterable)

        self.metadata.update(iterable.metadata)

    def _dict_to_token_and_set_defaults(self, token: T.Union[dict, Token]) -> Token:
        if not isinstance(token, Token):
            token = Token(token)

        if self.default_fields:
            for field in self.default_fields:
                if field not in token:
                    token[field] = "_"

        return token

    def append(self, token: T.Union[dict, Token]) -> None:
        token = self._dict_to_token_and_set_defaults(token)
        super(TokenList, self).append(token)

    def insert(self, i: int, token: T.Union[dict, Token]) -> None:
        token = self._dict_to_token_and_set_defaults(token)
        super(TokenList, self).insert(i, token)

    @T.overload
    def __setitem__(self, key: 'SupportsIndex', tokens: T.Union[dict, Token]) -> None: ...  # noqa, pragma: no cover

    @T.overload
    def __setitem__(self, key: slice, tokens: T.Union[T.Iterable[T.Union[dict, Token]], 'TokenList']) -> None: ...  # noqa, pragma: no cover

    def __setitem__(self, key, tokens):  # noqa: F811
        if isinstance(key, slice):
            tokens = [self._dict_to_token_and_set_defaults(token) for token in tokens]
            super(TokenList, self).__setitem__(key, tokens)
        else:
            token = tokens
            token = self._dict_to_token_and_set_defaults(token)
{'fraction_non_alphanumeric': 0.12614155251141554, 'fraction_numerical': 0.001141552511415525, 'mean_word_length': 2.4724534284581847, 'pattern_counts': {'":': 8, '<': 7, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '571660', 'n_tokens_mistral': 2726, 'n_tokens_neox': 2639, 'n_words': 694}