wesnoth/data/tools/wmlunits

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

Run without arguments to see usage.
"""

# Makes things faster on 32-bit systems
try: import psyco; psyco.full()
except ImportError: pass
  
import sys, os, glob, shutil, urllib2, argparse, traceback
import subprocess, yaml
import multiprocessing, Queue

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

TIMEOUT = 15

def copy_images():
    print("Recolorizing pictures.")
    image_collector.copy_and_color_images(options.output)
    shutil.copy2(os.path.join(image_collector.datadir,
        "data/tools/unit_tree/style.css"), options.output)
    shutil.copy2(os.path.join(image_collector.datadir,
        "data/tools/unit_tree/menu.js"), options.output)
    for grab in [
        "http://www.wesnoth.org/mw/skins/glamdrol/headerbg.jpg",
        "http://www.wesnoth.org/mw/skins/glamdrol/wesnoth-logo.jpg",
        "http://www.wesnoth.org/mw/skins/glamdrol/navbg.png"]:
        local = os.path.join(options.output, grab[grab.rfind("/") + 1:])
        if not os.path.exists(local):
            print "Fetching", grab
            url = urllib2.urlopen(grab)
            file(local, "w").write(url.read())

def shell(com):
    #print(com)
    p = subprocess.Popen(com, stdout = subprocess.PIPE,
        stderr = subprocess.PIPE, shell = True)
    out, err = p.communicate()
    #if out: sys.stdout.write(out)
    #if err: sys.stdout.write(err)

    return p.returncode

def shell_out(com):
    p = subprocess.Popen(com,
        stdout = subprocess.PIPE, stderr = subprocess.PIPE)
    out, err = p.communicate()
    return out

def bash(name):
    return "'" + name.replace("'", "'\\''") + "'"

def move(f, t, name):

    if os.path.exists(f + "/" + name + ".cfg"):
        com = "mv " + f + "/" + bash(name + ".cfg") + " " + t + "/"
        shell(com)

    com = "mv " + f + "/" + bash(name) + " " + t + "/"
    return shell(com)
   
_info = {}
def get_info(addon):
    global _info
    if addon in _info:
        return _info[addon]
    _info[addon] = None
    try:
        path = options.addons + "/" + addon + "/_info.cfg"
        if os.path.exists(path):
            parser = wmlparser2.Parser(options.wesnoth, options.config_dir,
                options.data_dir, no_preprocess = False)
            parser.parse_file(path)
            _info[addon] = parser
    except wmlparser2.WMLError as e:
        print(e)
    return _info[addon]

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
        if c:
            info[-1]["name"] = c.get_text_val("name")
        else:
            info[-1]["name"] = name
        info[-1]["units"] = "?"
        info[-1]["translations"] = {}
        
        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:
                info[-1]["translations"][isocode] = c.get_text_val("name",
                    translation = translate)
            else:
                info[-1]["translations"][isocode] = translate(name, domain)

    def list_eras(addon):
        eras = local.wesnoth.parser.get_all(tag = "era")
        if addon != "mainline":
            eras = [x for x in eras if not x.get_text_val("id") in mainline_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(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
            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(options.output + "/mainline")
    except OSError: pass

    try:
        batchlist = yaml.load(open(options.list))
    except IOError:
        batchlist = []
    
    def search(name):
        for info in batchlist:
            if info and info["name"] == name: return info
        batchlist.append({})
        batchlist[-1]["name"] = name
        return batchlist[-1]

    print("mainline")

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

    parse("{core}{multiplayer/eras.cfg}", "SKIP_CORE")
    info["eras"] = list_eras("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}", "SKIP_CORE")
        info["campaigns"] += list_campaigns("mainline")

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

    _deps = [{}]
    def get_dependencies(addon):
        if addon in _deps[0]:
            return _deps[0][addon]
        _deps[0][addon] = []
        try:
            info = get_info(addon).get_all(tag = "info")[0]
            row = info.get_text_val("dependencies")
            if row:
                deps1 = row.split(",")
            else:
                deps = []
            for d in deps1:
                if d in addons:
                    _deps[0][addon].append(d)
                else:
                    print("Missing dependency for " + addon + ": " + d)
        except Exception as e:
            pass
        return _deps[0][addon]

    sorted = []
    unsorted = addons[:]
    while unsorted:
        n = 0
        print(len(sorted))
        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
    addons = sorted

    for i, addon in enumerate(addons):
        if not os.path.isdir(options.addons + "/" + addon): continue
        sys.stdout.write("%4d/%4d " % (1 + i, len(addons)) + addon + " ... ")
        sys.stdout.flush()
        d = options.output + "/" + addon
        logname = d + "/error.log"
        try: os.makedirs(d)
        except OSError: pass
        version = get_version(addon)
        move(options.addons, options.config_dir + "/data/add-ons", addon)
        try:
            info = search(addon)

            if info.get("version", "") == version and info.get("parsed", False) == True:
                sys.stdout.write("up to date\n")
                continue
            info["parsed"] = False
            
            parse("{core}{multiplayer}{~add-ons}", "MULTIPLAYER,SKIP_CORE")
            info["eras"] = list_eras(addon)
            info["campaigns"] = list_campaigns(addon)
            info["version"] = version
            sys.stdout.write("ok\n")
        except wmlparser2.WMLError as e:
            ef = open(logname, "w")
            ef.write("<PARSE ERROR>\n")
            ef.write(str(e))
            ef.write("</PARSE ERROR>\n")
            ef.close()
            sys.stdout.write("failed\n")
        except Queue.Empty as e:
            ef = open(logname, "w")
            ef.write("<TIMEOUT ERROR>\n")
            ef.write("Failed to parse the WML within " + str(TIMEOUT) + " seconds.")
            ef.write("</TIMEOUT ERROR>\n")
            ef.close()
            sys.stdout.write("failed\n")
        except Exception as e:
            ef = open(logname, "w")
            ef.write("<INTERNAL ERROR>\n")
            ef.write(str(e))
            ef.write("</INTERNAL ERROR>\n")
            ef.close()
            sys.stdout.write("failed\n")

        finally:
            move(options.config_dir + "/data/add-ons", options.addons, addon)
        
    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(addon + ": " + cid + " " + 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}", "NORMAL")
    wesnoth.add_units("mainline")
    
    if define == "MULTIPLAYER":
        wesnoth.parser.parse_text("{core}{multiplayer}{~add-ons}", "MULTIPLAYER,SKIP_CORE")
        wesnoth.add_units(cid)
    else:
        if addon == "mainline":
            if cid != "mainline":
                wesnoth.parser.parse_text("{core}{campaigns}", "SKIP_CORE,NORMAL," + define)
                wesnoth.add_units(cid)
        else:
            wesnoth.parser.parse_text("{core}{~add-ons}", "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 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 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.load(open(options.batch))
    for addon in batchlist:
        name = addon["name"]
        
        if not options.reparse and addon.get("parsed", False) == True: continue
        
        if name == "mainline":
            worked = True
        else:
            worked = (move(options.addons, options.config_dir + "/data/add-ons", name) == 0)
        d = options.output + "/" + name
        try: os.makedirs(d)
        except OSError: pass
        logname = 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 == None: cid = campaign["define"]
                if cid == None: cid = name
                n = process_campaign_or_era(name, cid, campaign["define"], batchlist)
                campaign["units"] = n
            
        except wmlparser2.WMLError as e:
            ef = open(logname, "a")
            ef.write("<WML ERROR>\n")
            ef.write(str(e))
            ef.write("</WML ERROR>\n")
            ef.close()
            print("    " + name + " failed")
        except Exception as e:
            traceback.print_exc()
            print("    " + name + " failed")
            ef = open(logname, "a")
            ef.write("<INTERNAL ERROR>\n")
            ef.write("please report as bug")
            ef.write("</INTERNAL ERROR>\n")
            ef.close()
        finally:
            if name != "mainline":
                move(options.config_dir + "/data/add-ons", options.addons, name)
                
        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 = 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 != 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, "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 --config-path).")
    ap.add_argument("-D", "--data-dir",
        help="Specify the wesnoth data dir (wesnoth --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("-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")
        op.print_help()
        sys.exit(-1)

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

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

    if not options.data_dir:
        options.data_dir = shell_out([options.wesnoth, "--path"]).strip()
        print("Using " + options.data_dir + " as data dir.")

    if not options.config_dir:
        options.config_dir = shell_out([options.wesnoth, "--config-path"]).strip()
        print("Using " + options.config_dir + " as config dir.")

    if not options.transdir:
        options.transdir = os.getcwd()
    
    if options.wiki:
        wiki_output.main()
        sys.exit(0)

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

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

        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)