<commit_before>import os.path
import re

from .trie import Trie
from .utils import hexstrings2int


COLL_ELEMENT_PATTERN = r"""
    \[
    (\*|\.)
    ([0-9A-Fa-f]{4})
    \.
    ([0-9A-Fa-f]{4})
    \.
    ([0-9A-Fa-f]{4})
    (?:\.([0-9A-Fa-f]{4}))?
\]
"""


class Collator:

    def __init__(self, filename=None):

        if filename is None:
            filename = os.path.join(os.path.dirname(__file__), "allkeys.txt")
        self.table = Trie()
        self.load(filename)

    def load(self, filename):
        with open(filename) as keys_file:
            for line in keys_file:
                line = line.split("#")[0].split("%")[0].strip()

                if not line:
                    continue

                if line.startswith("@version"):
                    pass
                else:
                    a, b = line.split(";")
                    char_list = hexstrings2int(a.split())
                    coll_elements = []
                    for x in re.finditer(COLL_ELEMENT_PATTERN, b.strip(), re.X):
                        alt, weight1, weight2, weight3, weight4 = x.groups()
                        weights = [weight1, weight2, weight3]
                        if weight4:
                            weights.append(weight4)
                        coll_elements.append(hexstrings2int(weights))
                    self.table.add(char_list, coll_elements)

    def collation_elements(self, string):
        collation_elements = []

        lookup_key = [ord(ch) for ch in string]
        while lookup_key:
            value, lookup_key = self.table.find_prefix(lookup_key)
            if not value:
                # Calculate implicit weighting for CJK Ideographs
                # http://www.unicode.org/reports/tr10/#Implicit_Weights
                key = lookup_key[0]
                value = [
                    (0xFB40 + (key >> 15), 0x0020, 0x0002, 0x0001),
                    ((key & 0x7FFF) | 0x8000, 0x0000, 0x0000, 0x0000)
                ]
                lookup_key = lookup_key[1:]
            collation_elements.extend(value)

        return collation_elements

    def sort_key_from_collation_elements(self, collation_elements):
        sort_key = []

        for level in range(4):
            if level:
                sort_key.append(0)  # level separator
            for element in collation_elements:
                ce_l = element[level]
                if ce_l:
                    sort_key.append(ce_l)

        return tuple(sort_key)

    def sort_key(self, string):

        collation_elements = self.collation_elements(string)
        return self.sort_key_from_collation_elements(collation_elements)<commit_msg>handle case where level 4 weight isn't there<commit_after>import os.path
import re

from .trie import Trie
from .utils import hexstrings2int


COLL_ELEMENT_PATTERN = r"""
    \[
    (\*|\.)
    ([0-9A-Fa-f]{4})
    \.
    ([0-9A-Fa-f]{4})
    \.
    ([0-9A-Fa-f]{4})
    (?:\.([0-9A-Fa-f]{4}))?
\]
"""


class Collator:

    def __init__(self, filename=None):

        if filename is None:
            filename = os.path.join(os.path.dirname(__file__), "allkeys.txt")
        self.table = Trie()
        self.load(filename)

    def load(self, filename):
        with open(filename) as keys_file:
            for line in keys_file:
                line = line.split("#")[0].split("%")[0].strip()

                if not line:
                    continue

                if line.startswith("@version"):
                    pass
                else:
                    a, b = line.split(";")
                    char_list = hexstrings2int(a.split())
                    coll_elements = []
                    for x in re.finditer(COLL_ELEMENT_PATTERN, b.strip(), re.X):
                        alt, weight1, weight2, weight3, weight4 = x.groups()
                        weights = [weight1, weight2, weight3]
                        if weight4:
                            weights.append(weight4)
                        coll_elements.append(hexstrings2int(weights))
                    self.table.add(char_list, coll_elements)

    def collation_elements(self, string):
        collation_elements = []

        lookup_key = [ord(ch) for ch in string]
        while lookup_key:
            value, lookup_key = self.table.find_prefix(lookup_key)
            if not value:
                # Calculate implicit weighting for CJK Ideographs
                # http://www.unicode.org/reports/tr10/#Implicit_Weights
                key = lookup_key[0]
                value = [
                    (0xFB40 + (key >> 15), 0x0020, 0x0002, 0x0001),
                    ((key & 0x7FFF) | 0x8000, 0x0000, 0x0000, 0x0000)
                ]
                lookup_key = lookup_key[1:]
            collation_elements.extend(value)

        return collation_elements

    def sort_key_from_collation_elements(self, collation_elements):
        sort_key = []

        for level in range(4):
            if level:
                sort_key.append(0)  # level separator
            for element in collation_elements:
                if len(element) > level:
                    ce_l = element[level]
                    if ce_l:
                        sort_key.append(ce_l)

        return tuple(sort_key)

    def sort_key(self, string):

        collation_elements = self.collation_elements(string)
        return self.sort_key_from_collation_elements(collation_elements)
{'fraction_non_alphanumeric': 0.10740604274134119, 'fraction_numerical': 0.027450257921886514, 'mean_word_length': 1.5488262910798123, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '2310346', 'n_tokens_mistral': 1744, 'n_tokens_neox': 1549, 'n_words': 325}