acks-tools/acks/quest/commands.py

import click
from flask.cli import AppGroup

from ..models import db


npc_cli = AppGroup('npc')

@npc_cli.command('populate')
def populate_npc_database():
    import csv
    from .models import (
        CharacterClass,
        ClassLevelProgression,
        EquipmentArmour,
        EquipmentRangedWeapon,
        EquipmentMeleeWeapon,
        Spell
    )

    def load_csv_data(file_name, cls):
        rows = []
        with open('acks/npc/data/{}'.format(file_name), newline='') as data:
            reader = csv.DictReader(data)
            for row in reader:
                rows.append(cls(**row))
        return rows

    # Character Classes
    db.session.bulk_save_objects(load_csv_data('default_classes.csv', CharacterClass))

    # Equipment Armour
    db.session.bulk_save_objects(load_csv_data('default_armours.csv', EquipmentArmour))

    # Ranged Weapons
    db.session.bulk_save_objects(load_csv_data('default_ranged.csv', EquipmentRangedWeapon))

    # Melee Weapons
    melee_weps = load_csv_data('default_melee.csv', EquipmentMeleeWeapon)
    for wep in melee_weps:
        wep.two_handed = (wep.two_handed == 'True')
    db.session.bulk_save_objects(melee_weps)

    # Level Progressions
    progressions = load_csv_data('default_progression.csv', ClassLevelProgression)
    classes = {c.name: c.id for c in CharacterClass.query.all()}
    for prog in progressions:
        prog.guild_id = classes[prog.guild_id]
    db.session.bulk_save_objects(progressions)

    # Spells
    spells = load_csv_data('default_spells.csv', Spell)
    for spell in spells:
        if spell.arcane == '':
            spell.arcane = 0
        if spell.divine == '':
            spell.divine = 0
        spell.description = spell.description.strip()
    db.session.bulk_save_objects(spells)

    db.session.commit()