EDMarketConnector/collate.py

#!/usr/bin/env python3
"""Collate lists of seen commodities, modules and ships from dumps of the Companion API output."""

import csv
import json
import os
import sys
from os.path import isfile
from traceback import print_exc

import companion
import outfitting
from edmc_data import companion_category_map, ship_name_map


def __make_backup(file_name: str, suffix: str = '.bak') -> None:
    """
    Rename the given file to $file.bak, removing any existing $file.bak. Assumes $file exists on disk.

    :param file_name: The name of the file to make a backup of
    :param suffix: The suffix to use for backup files (default '.bak')
    """
    backup_name = file_name + suffix

    if isfile(backup_name):
        os.unlink(backup_name)

    os.rename(file_name, backup_name)


def addcommodities(data) -> None:  # noqa: CCR001
    """
    Keep a summary of commodities found using in-game names.

    Assumes that the commodity data has already been 'fixed up'
    :param data: - Fixed up commodity data.
    """
    if not data['lastStarport'].get('commodities'):
        return

    commodityfile = 'commodity.csv'
    commodities = {}

    # slurp existing
    if isfile(commodityfile):
        with open(commodityfile) as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                commodities[int(row['id'])] = row  # index by int for easier lookup and sorting

    size_pre = len(commodities)

    for commodity in data['lastStarport'].get('commodities'):
        key = int(commodity['id'])
        new = {
            'id':        commodity['id'],
            'symbol':    commodity['name'],
            'category':  companion_category_map.get(commodity['categoryname']) or commodity['categoryname'],
            'name':      commodity.get('locName') or 'Limpets',
        }

        old = commodities.get(key)

        if old and companion_category_map.get(commodity['categoryname'], True):
            if new['symbol'] != old['symbol'] or new['name'] != old['name']:
                raise ValueError(f'{key}: {new!r} != {old!r}')

        commodities[key] = new

    if not len(commodities) > size_pre:
        return

    if isfile(commodityfile):
        __make_backup(commodityfile)

    with open(commodityfile, 'w', newline='\n') as csvfile:
        writer = csv.DictWriter(csvfile, ['id', 'symbol', 'category', 'name'])
        writer.writeheader()

        for key in sorted(commodities):
            writer.writerow(commodities[key])

    print(f'Added {len(commodities) - size_pre} new commodities')


def addmodules(data):  # noqa: C901, CCR001
    """Keep a summary of modules found."""
    if not data['lastStarport'].get('modules'):
        return

    outfile = 'outfitting.csv'
    modules = {}
    fields = ('id', 'symbol', 'category', 'name', 'mount', 'guidance', 'ship', 'class', 'rating', 'entitlement')

    # slurp existing
    if isfile(outfile):
        with open(outfile) as csvfile:
            reader = csv.DictReader(csvfile, restval='')
            for row in reader:
                modules[int(row['id'])] = row  # index by int for easier lookup and sorting

    size_pre = len(modules)

    for key, module in data['lastStarport'].get('modules').items():
        # sanity check
        key = int(key)
        if key != module.get('id'):
            raise ValueError(f'id: {key} != {module["id"]}')

        try:
            new = outfitting.lookup(module, ship_name_map, True)

        except Exception:
            print(f'{module["id"]}, {module["name"]}:')
            print_exc(0)
            new = None

        if new:
            old = modules.get(key)
            if old:
                # check consistency with existing data
                for thing in fields:
                    if not old.get(thing) and new.get(thing):
                        size_pre -= 1

                    elif str(new.get(thing, '')) != old.get(thing):
                        raise ValueError(f'{key}: {thing} {new.get(thing)!r}!={old.get(thing)!r}')

            modules[key] = new

    if not len(modules) > size_pre:
        return

    if isfile(outfile):
        __make_backup(outfile)

    with open(outfile, 'w', newline='\n') as csvfile:
        writer = csv.DictWriter(csvfile, fields, extrasaction='ignore')
        writer.writeheader()

        for key in sorted(modules):
            writer.writerow(modules[key])

    print(f'Added {len(modules) - size_pre} new modules')


# keep a summary of ships found
def addships(data):
    if not data['lastStarport'].get('ships'):
        return

    shipfile = 'shipyard.csv'
    ships = {}
    fields = ('id', 'symbol', 'name')

    # slurp existing
    if isfile(shipfile):
        with open(shipfile) as csvfile:
            reader = csv.DictReader(csvfile, restval='')
            for row in reader:
                ships[int(row['id'])] = row  # index by int for easier lookup and sorting

    size_pre = len(ships)

    data_ships = data['lastStarport']['ships']
    for ship in tuple(data_ships.get('shipyard_list', {}).values()) + data_ships.get('unavailable_list'):
        # sanity check
        key = int(ship['id'])
        new = {'id': key, 'symbol': ship['name'], 'name': ship_name_map.get(ship['name'].lower())}
        if new:
            old = ships.get(key)
            if old:
                # check consistency with existing data
                for thing in fields:
                    if not old.get(thing) and new.get(thing):
                        ships[key] = new
                        size_pre -= 1

                    elif str(new.get(thing, '')) != old.get(thing):
                        raise ValueError(f'{key}: {thing} {new.get(thing)!r} != {old.get(thing)!r}')

            ships[key] = new

    if not len(ships) > size_pre:
        return

    if isfile(shipfile):
        __make_backup(shipfile)

    with open(shipfile, 'w', newline='\n') as csvfile:
        writer = csv.DictWriter(csvfile, ['id', 'symbol', 'name'])
        writer.writeheader()

        for key in sorted(ships):
            writer.writerow(ships[key])

    print(f'Added {len(ships) - size_pre} new ships')


if __name__ == "__main__":
    if len(sys.argv) <= 1:
        print('Usage: collate.py [dump.json]')
        sys.exit()

    # read from dumped json file(s)
    session = companion.Session()
    for file_name in sys.argv[1:]:
        data = None
        with open(file_name) as f:
            print(file_name)
            data = json.load(f)

        if not data['commander'].get('docked'):
            print('Not docked!')
            continue

        elif not data.get('lastStarport'):
            print('No starport!')
            continue

        if data['lastStarport'].get('commodities'):
            addcommodities(data)

        else:
            print('No market')

        if data['lastStarport'].get('modules'):
            addmodules(data)

        else:
            print('No outfitting')

        if data['lastStarport'].get('ships'):
            addships(data)

        else:
            print('No shipyard')