calefaction/scripts/import_sde.py

#!/usr/bin/env python3
# -*- coding: utf-8  -*-

import gzip
from pathlib import Path
import shutil
import sys

import yaml

_REGION = 3
_CONSTELLATION = 4
_SOLAR_SYSTEM = 5

def _load_yaml(filename):
    with filename.open("rb") as fp:
        return yaml.load(fp, Loader=yaml.CLoader)

def _verify_typeids(sde_dir):
    print("Verifying typeIDs... ", end="", flush=True)

    filename = sde_dir / "fsd" / "typeIDs.yaml"
    with filename.open("rb") as fp:
        data = yaml.load(fp.read(1024 * 16), Loader=yaml.CLoader)

    assert data[_REGION]["groupID"] == _REGION
    assert data[_REGION]["name"]["en"] == "Region"
    assert data[_CONSTELLATION]["groupID"] == _CONSTELLATION
    assert data[_CONSTELLATION]["name"]["en"] == "Constellation"
    assert data[_SOLAR_SYSTEM]["groupID"] == _SOLAR_SYSTEM
    assert data[_SOLAR_SYSTEM]["name"]["en"] == "Solar System"

    print("done.")

def _load_ids(sde_dir):
    print("Loading itemIDs... ", end="", flush=True)

    data = _load_yaml(sde_dir / "bsd" / "invItems.yaml")

    ids = {_REGION: [], _CONSTELLATION: [], _SOLAR_SYSTEM: []}
    for entry in data:
        if entry["typeID"] in ids:
            ids[entry["typeID"]].append(entry["itemID"])

    print("done.")
    return ids

def _load_names(sde_dir):
    print("Loading itemNames... ", end="", flush=True)

    data = _load_yaml(sde_dir / "bsd" / "invNames.yaml")

    names = {}
    for entry in data:
        name = entry["itemName"]
        assert isinstance(name, str)
        names[entry["itemID"]] = name

    print("done.")
    return names

def _build_galaxy_skeleton(ids, names):
    print("Building galaxy skeleton... ", end="", flush=True)

    galaxy = {"regions": {}, "constellations": {}, "systems": {}}

    d = galaxy["regions"]
    for rid in ids[_REGION]:
        assert isinstance(rid, int)
        d[rid] = {
            "name": names[rid]
        }

    d = galaxy["constellations"]
    for cid in ids[_CONSTELLATION]:
        assert isinstance(cid, int)
        d[cid] = {
            "name": names[cid],
            "region": -1
        }

    d = galaxy["systems"]
    for sid in ids[_SOLAR_SYSTEM]:
        assert isinstance(sid, int)
        d[sid] = {
            "name": names[sid],
            "constellation": -1, "region": -1, "security": 0.0
        }

    print("done.")
    return galaxy

def _load_galaxy_associations(sde_dir, galaxy):
    print("Loading galaxy staticdata... ", end="", flush=True)

    univdir = sde_dir / "fsd" / "universe"
    for base in univdir.iterdir():
        if not base.is_dir():
            continue

        for region in base.iterdir():
            if not region.is_dir():
                continue
            rdata = _load_yaml(region / "region.staticdata")
            rid = rdata["regionID"]

            assert isinstance(rid, int)
            assert rid >= 0

            for constellation in region.iterdir():
                if not constellation.is_dir():
                    continue
                cdata = _load_yaml(constellation / "constellation.staticdata")
                cid = cdata["constellationID"]

                assert isinstance(cid, int)
                assert cid >= 0

                galaxy["constellations"][cid]["region"] = rid

                for system in constellation.iterdir():
                    if not system.is_dir():
                        continue
                    sdata = _load_yaml(system / "solarsystem.staticdata")
                    sid = sdata["solarSystemID"]
                    sec = sdata["security"]

                    assert isinstance(sid, int)
                    assert isinstance(sec, float)
                    assert sid >= 0
                    assert sec >= -1.0 and sec <= 1.0

                    galaxy["systems"][sid]["constellation"] = cid
                    galaxy["systems"][sid]["region"] = rid
                    galaxy["systems"][sid]["security"] = sec

    print("done.")

    for cid, constellation in galaxy["constellations"].items():
        if constellation["region"] < 0:
            print("[WARNING] Orphaned constellation: %d=%s" % (
                cid, constellation["name"]))

    for sid, system in galaxy["systems"].items():
        if system["region"] < 0 or system["constellation"] < 0:
            print("[WARNING] Orphaned system: %d=%s" % (sid, system["name"]))

def _dump_galaxy(out_dir, galaxy):
    print("Dumping galaxy... ", end="", flush=True)

    filename = out_dir / "galaxy.yml"
    with filename.open("w") as fp:
        fp.write(yaml.dump(galaxy, Dumper=yaml.CDumper))

    print("done.")

def _compress_galaxy(out_dir):
    print("Compressing galaxy... ", end="", flush=True)

    fn_src = out_dir / "galaxy.yml"
    fn_dst = out_dir / "galaxy.yml.gz"

    with fn_src.open("rb") as f_in:
        with gzip.open(str(fn_dst), "wb") as f_out:
            shutil.copyfileobj(f_in, f_out)

    print("done.")

def _cleanup(out_dir):
    print("Cleaning up... ", end="", flush=True)

    (out_dir / "galaxy.yml").unlink()

    print("done.")

def import_sde(sde_dir, out_dir):
    """Import the SDE unzipped at sde_dir to out_dir."""
    print("EVE Online static data import")
    print("- from: %s" % sde_dir)
    print("- to:   %s" % out_dir)

    _verify_typeids(sde_dir)
    ids = _load_ids(sde_dir)
    print("Counts: regions=%d, constellations=%d, systems=%d" % (
        len(ids[_REGION]), len(ids[_CONSTELLATION]), len(ids[_SOLAR_SYSTEM])))
    names = _load_names(sde_dir)
    galaxy = _build_galaxy_skeleton(ids, names)
    del ids
    del names
    _load_galaxy_associations(sde_dir, galaxy)
    _dump_galaxy(out_dir, galaxy)
    del galaxy
    _compress_galaxy(out_dir)
    _cleanup(out_dir)

def main():
    if len(sys.argv) < 2:
        print("usage: %s <sde_directory>" % sys.argv[0])
        exit(1)

    sde_dir = Path(sys.argv[1]).resolve()
    out_dir = Path(__file__).resolve().parent.parent / "data" / "universe"
    import_sde(sde_dir, out_dir)

if __name__ == "__main__":
    main()