<commit_before>from lib import contract, base

from . import file_loader


class AttackTypes(base.BaseEnum):
    pass


class Attack(base.BaseClass):
    @contract.self_accepts(str, int, int)
    def __init__(self, name, power, attackType):
        self.name = name
        self.power = power
        self.attackType = attackType


@contract.returns(AttackTypes)
def load_attack_types():
    attack_types = AttackTypes()
    attack_list = file_loader.read_and_parse_json('attack_types')[0]
    for enumeration, attack_type in enumerate(attack_list):
        attack_types[str(attack_type)] = enumeration
    return attack_types


@contract.accepts(AttackTypes)
def load_attacks(attack_types):
    args = {}
    for attack_ in file_loader.read_and_parse_json('attacks'):
        name = str(attack_['name'])
        power = attack_['power']
        attack_type = attack_types[str(attack_['attack_type'])]
        args[name] = [name, power, attack_type]

    @contract.accepts(str)
    @contract.returns(Attack)
    def attack_getter(name):
        return Attack(*args[name])
    return attack_getter, args.keys()
<commit_msg>Use load_enum for attack types<commit_after>from lib import contract, base

from . import file_loader


class AttackTypes(base.BaseEnum):
    pass


class Attack(base.BaseClass):
    @contract.self_accepts(str, int, int)
    def __init__(self, name, power, attackType):
        self.name = name
        self.power = power
        self.attackType = attackType


@contract.returns(AttackTypes)
def load_attack_types():
    return AttackTypes(file_loader.load_enum('attack_types'))


@contract.accepts(AttackTypes)
def load_attacks(attack_types):
    args = {}
    for attack_ in file_loader.read_and_parse_json('attacks'):
        name = str(attack_['name'])
        power = attack_['power']
        attack_type = attack_types[str(attack_['attack_type'])]
        args[name] = [name, power, attack_type]

    @contract.accepts(str)
    @contract.returns(Attack)
    def attack_getter(name):
        return Attack(*args[name])
    return attack_getter, args.keys()

{'fraction_non_alphanumeric': 0.11084452975047984, 'fraction_numerical': 0.0004798464491362764, 'mean_word_length': 3.815242494226328, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '1316672', 'n_tokens_mistral': 683, 'n_tokens_neox': 645, 'n_words': 133}