wesnoth/data/tools/wmlunits

#!/usr/bin/env python3
# encoding: utf-8
"""
wmlunits -- tool to output information on all units in HTML

Run without arguments to see usage.
"""

import argparse
import multiprocessing
import os
import queue
import shutil
import subprocess
import sys
import traceback

# PyYAML needs to be downloaded and installed using a package manager, e.g. from PyPI using `pip install pyyaml`.
import yaml

import unit_tree.animations as animations
import unit_tree.helpers as helpers
import unit_tree.html_output as html_output
import unit_tree.overview
import unit_tree.wiki_output as wiki_output
import wesnoth.wmlparser3 as wmlparser3

TIMEOUT = 20

def copy_images():
    print("Recolorizing pictures...")
    image_collector.copy_and_color_images(options.output)
    shutil.copy2(os.path.join(os.path.dirname(os.path.realpath(__file__)), "unit_tree", "menu.js"),
                 options.output)

def shell_out(wesnoth_exe, com):
    base_com = [wesnoth_exe, "--nobanner"] + com
    cli_opt_matrix = [[], ["--no-log-to-file"], ["--wnoconsole", "--wnoredirect"]]
    for cli_opts in cli_opt_matrix:
        com = base_com + cli_opts
        p = subprocess.run(com, capture_output=True, text=True)
        if p.returncode == 0:
            return os.path.normpath(p.stdout.strip())
    return ""

def symlink(f, t, name):
    if os.path.exists(os.path.join(f, name + ".cfg")):
        os.symlink(os.path.join(f, name + ".cfg"), os.path.join(t, name + ".cfg"))
    os.symlink(os.path.join(f, name), os.path.join(t, name))

def rm_symlink(f, name):
    if os.path.exists(os.path.join(f, name + ".cfg")):
        os.remove(os.path.join(f, name + ".cfg"))
    os.remove(os.path.join(f, name))

_info = {}
def get_info(addon):
    global _info
    if addon in _info:
        return _info[addon]
    _info[addon] = None
    try:
        path = os.path.join(options.addons, addon, "_info.cfg")
        if os.path.exists(path):
            parser = wmlparser3.Parser(options.wesnoth,
                                       options.config_dir,
                                       options.data_dir)
            parser.parse_file(path)
            _info[addon] = parser
        else:
            print("Cannot find " + path)
    except wmlparser3.WMLError as e:
        print(e)
    return _info[addon]

_deps = {}
global_addons = set()
def get_dependencies(addon):
    global _deps
    global global_addons
    if addon in _deps:
        return _deps[addon]
    _deps[addon] = []
    try:
        info = get_info(addon).get_all(tag="info")[0]
        row = info.get_text_val("dependencies")
        if row:
            deps1 = row.split(",")
        else:
            deps1 = []
        for d in deps1:
            if d == addon:
                print("Addon " + addon + " depends on itself.")
                continue
            if d in global_addons:
                _deps[addon].append(d)
            else:
                print("Missing dependency for " + addon + ": " + d)
    except Exception as e:
        print(e)
    return _deps[addon]

def set_dependencies(addon, depends_on):
    _deps[addon] = depends_on

def get_all_dependencies(addon):
    result = []
    check = get_dependencies(addon)[:]
    while check:
        d = check.pop()
        if d == addon or d in result:
            continue
        result.append(d)
        check += get_dependencies(d)
    return result

def sorted_by_dependencies(addons):
    sorted = []
    unsorted = addons[:]
    while unsorted:
        n = 0
        for addon in unsorted:
            for d in get_dependencies(addon):
                if d not in sorted:
                    break
            else:
                sorted.append(addon)
                unsorted.remove(addon)
                n += 1
                continue
        if n == 0:
            print("Cannot sort dependencies for these addons: " + str(unsorted))
            sorted += unsorted
            break
    return sorted

def search(batchlist, name):
    for info in batchlist:
        if info and info["name"] == name:
            return info
    batchlist.append({})
    batchlist[-1]["name"] = name
    return batchlist[-1]

def list_contents():
    class Empty:
        pass

    local = Empty()
    mainline_eras = set()
    filename = options.list

    def append(info, id, define, c=None, name=None, domain=None):
        info.append({})
        info[-1]["id"] = id
        info[-1]["define"] = define
        info[-1]["name"] = c.get_text_val("name") if c else name
        info[-1]["units"] = "?"
        info[-1]["translations"] = {}

        for isocode in languages:
            translation = html_output.Translation(options.transdir, isocode)
            def translate(string, domain):
                return translation.translate(string, domain)
            if c:
                t = c.get_text_val("name", translation=translate)
            else:
                t = translate(name, domain)
            if t != info[-1]["name"]:
                info[-1]["translations"][isocode] = t

    def get_dependency_eras(batchlist, addon):
        dependency_eras = list(mainline_eras)
        for d in get_all_dependencies(addon):
            dinfo = search(batchlist, d)
            for era in dinfo.get("eras", []):
                dependency_eras.append(era["id"])
        return dependency_eras

    def list_eras(batchlist, addon):
        eras = local.wesnoth.parser.get_all(tag="era")
        if addon != "mainline":
            dependency_eras = get_dependency_eras(batchlist, addon)
            eras = [x for x in eras if not x.get_text_val("id") in dependency_eras]
        info = []
        for era in eras:
            eid = era.get_text_val("id")
            if addon == "mainline":
                mainline_eras.add(eid)
            append(info, eid, "MULTIPLAYER", c=era)

        return info

    def list_campaigns(batchlist, addon):
        campaigns = local.wesnoth.parser.get_all(tag="campaign")
        info = []

        for campaign in campaigns:
            cid = campaign.get_text_val("id")
            d = campaign.get_text_val("define")
            d2 = campaign.get_text_val("extra_defines")
            if d2:
                d += "," + d2
            d3 = campaign.get_text_val("difficulties")
            if d3:
                d += "," + d3.split(",")[0]
            else:
                difficulty = campaign.get_all(tag="difficulty")
                if difficulty:
                    d3 = difficulty[0].get_text_val("define")
                    if d3:
                        d += "," + d3

            append(info, cid, d, c=campaign)

        return info

    def parse(wml, defines):
        def f(options, wml, defines, q):
            local.wesnoth = helpers.WesnothList(
                options.wesnoth,
                options.config_dir,
                options.data_dir,
                options.transdir)
            #print("remote", local.wesnoth)
            try:
                local.wesnoth.parser.parse_text(wml, defines)
                q.put(("ok", local.wesnoth))
            except Exception as e:
                q.put(("e", e))

        q = multiprocessing.Queue()
        p = multiprocessing.Process(target=f, args=(options, wml, defines, q))
        p.start()
        try:
            s, local.wesnoth = q.get(timeout=TIMEOUT)
        except queue.Empty:
            p.terminate()
            raise
        #print("local", s, local.wesnoth)
        p.join()
        if s == "e":
            remote_exception = local.wesnoth
            raise remote_exception

    def get_version(addon):
        parser = get_info(addon)
        if parser:
            for info in parser.get_all(tag="info"):
                return info.get_text_val("version") + "*" + info.get_text_val("uploads")

    try:
        os.makedirs(os.path.join(options.output, "mainline"))
    except OSError:
        pass

    try:
        batchlist = yaml.safe_load(open(options.list))
    except IOError:
        batchlist = []

    print("mainline")

    info = search(batchlist, "mainline")
    info["version"] = "mainline"
    info["parsed"] = "false"

    parse("{core}{multiplayer/eras.cfg}", "__WMLUNITS__,SKIP_CORE")
    info["eras"] = list_eras(batchlist, "mainline")

    # Fake mainline campaign to have an overview of the mainline units
    info["campaigns"] = []
    append(info["campaigns"], "mainline", "", name="Units", domain="wesnoth-help")

    if not options.addons_only:
        parse("{core}{campaigns}", "__WMLUNITS__,SKIP_CORE")
        info["campaigns"] += list_campaigns(batchlist, "mainline")

    addons = []
    if options.addons:
        addons = os.listdir(options.addons)
    global global_addons
    global_addons = set(addons)

    # fill in the map for all dependencies
    for addon in addons:
        get_dependencies(addon)

    # this ensures that info about eras in dependant addons is available
    # already
    addons = sorted_by_dependencies(addons)

    for i, addon in enumerate(addons):
        if not os.path.isdir(os.path.join(options.addons, addon)):
            continue
        mem = "? KB"
        try:
            mem = dict([x.split("\t", 1) for x in open("/proc/self/status").read().split("\n") if x])["VmRSS:"]
        except:
            pass
        sys.stdout.write("%4d/%4d " % (1 + i, len(addons)) + addon + " [" + mem + "] ... ")
        sys.stdout.flush()
        d = os.path.join(options.output, addon)
        logname = os.path.join(d, "error.log")
        try:
            os.makedirs(d)
        except OSError:
            pass
        version = get_version(addon)
        symlink(options.addons, os.path.join(options.config_dir, "data", "add-ons"), addon)
        all_dependencies = get_all_dependencies(addon)
        for d in all_dependencies:
            symlink(options.addons, os.path.join(options.config_dir, "data", "add-ons"), d)
        try:
            info = search(batchlist, addon)
            if info.get("version", "") == version and info.get("parsed", False) is True:
                sys.stdout.write("up to date\n")
                continue
            info["parsed"] = False
            info["dependencies"] = all_dependencies
            parse("{themes}{core}{multiplayer}{multiplayer/eras.cfg}{~add-ons}", "__WMLUNITS__,MULTIPLAYER,SKIP_CORE")
            info["eras"] = list_eras(batchlist, addon)
            info["campaigns"] = list_campaigns(batchlist, addon)
            info["version"] = version
            print("ok")
        except wmlparser3.WMLError as e:
            with open(logname, "w") as ef:
                ef.write("<PARSE ERROR>\n")
                ef.write(str(e))
                ef.write("\n</PARSE ERROR>\n")
            print("failed")
        except queue.Empty as e:
            with open(logname, "w") as ef:
                ef.write("<TIMEOUT ERROR>\n")
                ef.write("Failed to parse the WML within " + str(TIMEOUT) + " seconds.")
                ef.write("\n</TIMEOUT ERROR>\n")
            print("failed")
        except Exception as e:
            with open(logname, "w") as ef:
                ef.write("<INTERNAL ERROR>\n")
                ef.write(repr(e))
                ef.write("\n</INTERNAL ERROR>\n")
            print("failed")
        finally:
            rm_symlink(os.path.join(options.config_dir, "data", "add-ons"), addon)
            for d in all_dependencies:
                rm_symlink(os.path.join(options.config_dir, "data", "add-ons"), d)

    yaml.safe_dump(batchlist, open(filename, "w"),
                   encoding="utf-8", default_flow_style=False)

def process_campaign_or_era(addon, cid, define, batchlist):
    n = 0

    print(str(addon) + ": " + str(cid) + " " + str(define))

    if not define:
        define = ""

    wesnoth = helpers.WesnothList(
        options.wesnoth,
        options.config_dir,
        options.data_dir,
        options.transdir)
    wesnoth.batchlist = batchlist
    wesnoth.cid = cid

    wesnoth.parser.parse_text("{core/units.cfg}", "__WMLUNITS__,NORMAL")
    wesnoth.add_units("mainline")

    if define == "MULTIPLAYER":
        wesnoth.parser.parse_text("{themes}{core}{multiplayer}{multiplayer/eras.cfg}{~add-ons}", "__WMLUNITS__,MULTIPLAYER,SKIP_CORE")
        wesnoth.add_units(cid)
    else:
        if addon == "mainline":
            if cid != "mainline":
                wesnoth.parser.parse_text("{core}{campaigns}", "__WMLUNITS__,SKIP_CORE,NORMAL," + define)
                wesnoth.add_units(cid)
        else:
            wesnoth.parser.parse_text("{themes}{core}{~add-ons}", "__WMLUNITS__,SKIP_CORE," + define)
            wesnoth.add_units(cid)

    if addon == "mainline" and cid == "mainline":
        write_animation_statistics(wesnoth)

    wesnoth.add_binary_paths(addon, image_collector)

    if define == "MULTIPLAYER":
        eras = wesnoth.parser.get_all(tag="era")
        for era in eras:
            wesnoth.add_era(era)
        wesnoth.find_unit_factions()
    else:
        campaigns = wesnoth.parser.get_all(tag="campaign")
        for campaign in campaigns:
            wesnoth.add_campaign(campaign)

    wesnoth.add_languages(languages)
    wesnoth.add_terrains()
    wesnoth.check_units()

    for isocode in languages:
        if addon != "mainline" and isocode != "en_US":
            continue
        if define == "MULTIPLAYER":
            for era in list(wesnoth.era_lookup.values()):
                if era.get_text_val("id") == cid:
                    n = html_output.generate_era_report(addon, isocode, era, wesnoth)
                    break
        else:
            if cid == "mainline":
                n = html_output.generate_campaign_report(addon, isocode, None, wesnoth)

            for campaign in list(wesnoth.campaign_lookup.values()):
                if campaign.get_text_val("id") == cid:
                    n = html_output.generate_campaign_report(addon, isocode, campaign, wesnoth)
                    break
        html_output.generate_single_unit_reports(addon, isocode, wesnoth)

    return n

def batch_process():
    batchlist = yaml.safe_load(open(options.batch))

    for addon in batchlist:
        name = addon["name"]
        set_dependencies(name, addon.get("dependencies", []))

    for addon in batchlist:
        name = addon["name"]

        if not options.reparse and addon.get("parsed", False) == True:
            continue

        if name == "mainline":
            worked = True
        else:
            try:
                symlink(options.addons, os.path.join(options.config_dir, "data", "add-ons"), name)
                worked = True
            except FileNotFoundError:
                worked = False
            for d in addon.get("dependencies", []):
                symlink(options.addons, os.path.join(options.config_dir, "data", "add-ons"), d)
        d = os.path.join(options.output, name)
        try:
            os.makedirs(d)
        except OSError:
            pass
        logname = os.path.join(d, "error.log")

        def err(mess):
            ef = open(logname, "a")
            ef.write(str(mess))
            ef.close()

        html_output.write_error = err

        try:
            if not worked:
                print(name + " not found")
                continue

            for era in addon.get("eras", []):
                eid = era["id"]
                n = process_campaign_or_era(name, eid, era["define"], batchlist)
                era["units"] = n

            for campaign in addon.get("campaigns", []):
                cid = campaign["id"]
                if cid is None:
                    cid = campaign["define"]
                if cid is None:
                    cid = name
                n = process_campaign_or_era(name, cid, campaign["define"], batchlist)
                campaign["units"] = n

        except wmlparser3.WMLError as e:
            print("    " + name + " failed")
            with open(logname, "a") as ef:
                ef.write("<WML ERROR>\n")
                ef.write(str(e))
                ef.write("\n</WML ERROR>\n")
        except (AttributeError, KeyError, TypeError) as e:
            # Common logic errors.
            traceback.print_exc()
            print("    " + name + " failed")
            with open(logname, "a") as ef:
                ef.write("<INTERNAL ERROR>\n")
                ef.write("please report as bug\n")
                # Show last few stack frames of traceback to make bug reports more useful.
                ef.write(traceback.format_exc(limit=-2))
                ef.write("\n</INTERNAL ERROR>\n")
        except Exception as e:
            traceback.print_exc()
            print("    " + name + " failed")
            with open(logname, "a") as ef:
                ef.write("<INTERNAL ERROR>\n")
                ef.write("please report as bug\n")
                ef.write(str(e))
                ef.write("\n</INTERNAL ERROR>\n")
        finally:
            if name != "mainline":
                rm_symlink(os.path.join(options.config_dir, "data", "add-ons"), name)
                for d in addon.get("dependencies", []):
                    rm_symlink(os.path.join(options.config_dir, "data", "add-ons"), d)

        addon["parsed"] = True

        yaml.safe_dump(batchlist, open(options.batch, "w"),
                       encoding="utf-8", default_flow_style=False)

        try:
            unit_tree.overview.write_addon_overview(
                os.path.join(options.output, name), addon)
        except Exception as e:
            pass

    html_output.html_postprocess_all(batchlist)

def write_unit_ids_UNUSED():
    # Write a list with all unit ids, just for fun.
    uids = list(wesnoth.unit_lookup.keys())
    def by_race(u1, u2):
        r = cmp(wesnoth.unit_lookup[u1].rid,
                wesnoth.unit_lookup[u2].rid)
        if r == 0:
            r = cmp(u1, u2)
        return r
    uids.sort(by_race)
    race = None
    f = MyFile(os.path.join(options.output, "uids.html"), "w")
    f.write("<html><body>")
    for uid in uids:
        u = wesnoth.unit_lookup[uid]
        if u.rid != race:
            if race is not None:
                f.write("</ul>")
            f.write("<p>%s</p>\n" % (u.rid,))
            f.write("<ul>")
            race = u.rid
        f.write("<li>%s</li>\n" % (uid, ))
    f.write("</ul>")
    f.write("</body></html>")
    f.close()

def write_animation_statistics(wesnoth):
    # Write animation statistics
    f = html_output.MyFile(os.path.join(options.output, "animations.html"), "w")
    animations.write_table(f, wesnoth)
    f.close()

if __name__ == '__main__':
    # We change the process name to "wmlunits"
    try:
        import ctypes
        libc = ctypes.CDLL("libc.so.6")
        libc.prctl(15, b"wmlunits", 0, 0, 0)
    except: # oh well...
        pass

    global options
    global image_collector

    ap = argparse.ArgumentParser()
    ap.add_argument("-C", "--config-dir",
                    help="Specify the user configuration dir (wesnoth --userdata-path).")
    ap.add_argument("-D", "--data-dir",
                    help="Specify the wesnoth data dir (wesnoth --data-path).")
    ap.add_argument("-l", "--language", default="all",
                    help="Specify a language to use. Else output is produced for all languages.")
    ap.add_argument("-o", "--output",
                    help="Specify the output directory.")
    ap.add_argument("-n", "--nocopy", action="store_true",
                    help="No copying of files. By default all images are copied to the output dir.")
    ap.add_argument("-w", "--wesnoth",
                    help="Specify the wesnoth executable to use. Whatever data " +
                    "and config paths that executable is configured for will be " +
                    "used to find game files and addons.")
    ap.add_argument("-t", "--transdir",
                    help="Specify the directory with gettext message catalogues. " +
                    "Defaults to ./translations.", default="translations")
    ap.add_argument("-m", "--magick", default="convert",
                    help="Command that executes the convert(1) program of the " +
                    "imagemagick(1) suite. Supports subcommands (e.g. `magick convert`).")
    ap.add_argument("-T", "--timeout",
                    help="Use the timeout iven in seconds when parsing WML, default is 20 " +
                    "seconds.")
    ap.add_argument("-r", "--reparse", action="store_true",
                    help="Reparse everything.")
    ap.add_argument("-a", "--addons",
                    help="Specify path to a folder with all addons. This should be " +
                    "outside the user config folder.")
    ap.add_argument("-L", "--list",
                    help="List available eras and campaigns.")
    ap.add_argument("-B", "--batch",
                    help="Batch process the given list.")
    ap.add_argument("-A", "--addons-only", action="store_true",
                    help="Do only process addons (for debugging).")
    ap.add_argument("-v", "--verbose", action="store_true")
    ap.add_argument("-W", "--wiki", action="store_true",
                    help="write wikified units list to stdout")
    options = ap.parse_args()

    html_output.options = options
    helpers.options = options
    unit_tree.overview.options = options
    wiki_output.options = options

    if not options.output and not options.wiki:
        sys.stderr.write("Need --output (or --wiki).\n")
        ap.print_help()
        sys.exit(-1)

    if options.output:
        options.output = os.path.expanduser(options.output)

    if options.timeout:
        TIMEOUT = int(options.timeout)

    if not options.wesnoth:
        options.wesnoth = "wesnoth"

    if not options.data_dir:
        options.data_dir = shell_out(options.wesnoth, ["--data-path"]).strip()
        if not options.data_dir:
            sys.stderr.write("Need --data-dir.\n")
            ap.print_help()
            sys.exit(-1)
        print("Using " + options.data_dir + " as data dir.")

    if not options.config_dir:
        options.config_dir = shell_out(options.wesnoth, ["--userdata-path"]).strip()
        if not options.config_dir:
            sys.stderr.write("Need --config-dir.\n")
            ap.print_help()
            sys.exit(-1)
        print("Using " + options.config_dir + " as config dir.")

    if not options.transdir:
        options.transdir = os.getcwd()

    options.data_dir = os.path.normpath(options.data_dir)
    options.config_dir = os.path.normpath(options.config_dir)
    options.transdir = os.path.normpath(options.transdir)
    if options.addons:
        options.addons = os.path.normpath(options.addons)

    if options.wiki:
        wiki_output.main()
        sys.exit(0)

    image_collector = helpers.ImageCollector(options.wesnoth,
                                             options.config_dir,
                                             options.data_dir,
                                             magick_exe=options.magick)
    html_output.image_collector = image_collector

    if options.language == "all":
        languages = []
        parser = wmlparser3.Parser(options.wesnoth,
                                   options.config_dir,
                                   options.data_dir)
        parser.parse_text("{languages}", "__WMLUNITS__")

        for locale in parser.get_all(tag="locale"):
            isocode = locale.get_text_val("locale")
            name = locale.get_text_val("name")
            if isocode == "ang_GB":
                continue
            languages.append(isocode)
        languages.sort()
    else:
        languages = options.language.split(",")

    if not options.list and not options.batch:
        sys.stderr.write("Need --list or --batch (or both).\n")
        sys.exit(-1)

    if options.output:
        # Generate output dir.
        if not os.path.isdir(options.output):
            os.mkdir(options.output)

    if options.list:
        list_contents()

    if options.batch:
        batch_process()

        unit_tree.overview.main(options.output)

        if not options.nocopy:
            copy_images()

    html_output.write_index(options.output)