wesnoth/data/tools/wmlunits

#!/usr/bin/env python
"""
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, re, glob, shutil, copy

import wesnoth.wmldata as wmldata
import wesnoth.wmlparser as wmlparser
import wesnoth.wmltools as wmltools

class ParserWithCoreMacros:
    def __init__(self, isocode):
        # Handle translations.
        self.translations = wmltools.Translations()
        def gettext(textdomain, x):
            return self.translations.get(textdomain, isocode, x, x)
        self.gettext = gettext

        # Create a new parser for the macros.
        parser = wmlparser.Parser(datadir)
        parser.do_preprocessor_logic = True
        parser.gettext = self.gettext
        
        # Parse core macros.
        parser.parse_text("{core/macros/}\n")
        parser.parse_top(None)
        self.core_macros = parser.macros
                
    def parse(self, text_to_parse):
        # Create the real parser.
        parser = wmlparser.Parser(datadir)
        parser.do_preprocessor_logic = True
        parser.gettext = self.gettext
        parser.macros = copy.copy(self.core_macros)
        
        # Suppress complaints about undefined terrain macros
        parser.set_macro_not_found_callback(lambda wmlparser, name, params:
        name.startswith("TERRAIN") or name == "DISABLE_TRANSITIONS")

        # Create a WML root element and parse the given text into it.
        WML = wmldata.DataSub("WML")
        parser.parse_text(text_to_parse)
        parser.parse_top(WML)
 
        return WML

class ImageCollector:
    def __init__(self):
        self.images = {}
    def add(self, campaign, path):
        self.images[(path, campaign)] = True
    def copy_and_color_images(self, target_path):
        for i, c in self.images.keys():

            if c == "mainline": bases = ["core/images"]
	    else: bases = ["campaigns/%s/images" % c, "campaigns/%s/images/units" % c, "core/images"]
	    for base in bases:
                ipath = os.path.join("data/%s" % base, i)
		if os.path.exists(ipath): break
            opath = os.path.join(target_path, "pics", c, i)
            try:
                os.makedirs(os.path.dirname(opath))
            except OSError:
                pass
            if os.path.exists(ipath):
                #shutil.copy2(ipath, opath)
                command = "data/tools/unit_tree/TeamColorizer %s %s" % (
                    ipath, opath)
                os.system(command)
            else:
                sys.stderr.write(
                    "Warning: Required image %s: \"%s\" does not exist.\n" % (c, i))

image_collector = ImageCollector()

class UnitList:
    def __init__(self, isocode):
        self.units_by_campaign = {}
        self.unit_lookup = {}
        self.race_lookup = {}
        self.terrain_lookup = {}
        self.movetype_lookup = {}
        self.parser = ParserWithCoreMacros(isocode)

    def add_terrains(self):
        WML = self.parser.parse("{core/terrain.cfg}\n")

        for terrain in WML.get_all("terrain"):
            tid = terrain.get_text_val("id")
            self.terrain_lookup[tid] = terrain

    def add(self, text_to_parse, campaign):
        "Collect all units in the specified namespace."

        WML = self.parser.parse(text_to_parse)

        # Collect unit data. First, we look for a [+units] section.
        units = WML.get_first("+units")
	    # If no [+units] section, assume it is inside a [campaign].
        if not units:
            campaign_wml = WML.get_first("campaign")
            if not campaign_wml:
                # FIXME: The tutorial has no [campaign], need to special case it
                # somehow. Currently, we ignore it.
                return 0
            # Now we get the define - strange, but seems to be how Wesnoth
            # works..
            define = campaign_wml.get_text_val("define")

            # Re-parse, this time with the define defined.
            WML = self.parser.parse(
                "#define %s\n#enddef\n%s" % (define, text_to_parse))
            # This time, it oughta work.
            units = WML.get_first("+units")
            if not units:
                # This campaign has no units. Nothing to do.
                return 0

        # Find all units types.
        newunits = units.get_all("unit_type")
        self.units_by_campaign[campaign] = []
        for unit in newunits:
            if unit.get_text_val("do_not_list", "no") == "no":
                uid = unit.get_text_val("id")
                self.units_by_campaign[campaign].append(unit)
                self.unit_lookup[uid] = unit
                unit.campaign = campaign

        # Find all races.
        newraces = units.get_all("race")
        for race in newraces:
            rid = race.get_text_val("id")
            self.race_lookup[rid] = race
        
        # Find all movetypes.
        newmovetypes = units.get_all("movetype")
        for movetype in newmovetypes:
            mtname = movetype.get_text_val("name")
            self.movetype_lookup[mtname] = movetype

        # Store the race and movetype of each unit for easier access later.
        for unit in newunits:
            race = self.get_unit_value(unit, "race")
            try: unit.race = self.race_lookup[race]
            except KeyError:
                unit.race = None
                sys.stderr.write("Warning: No race \"%s\" found (%s).\n" % (
                    race, unit.get_text_val("id")))
            movetype = self.get_unit_value(unit, "movement_type")
            try: unit.movetype = self.movetype_lookup[movetype]
            except KeyError: unit.movetype = None

        return len(newunits)

    def get_base_unit(self, unit):
        b = unit.get_first("base_unit")
        if b:
            buid = b.get_text_val("id")
            try: baseunit = self.unit_lookup[buid]
            except KeyError:
                sys.stderr.write(
                    "Warning: No baseunit \"%s\" for \"%s\".\n" % (
                    buid, unit.get_text_val("id")))
                return None
            return baseunit
        return None

    def get_unit_value(self, unit, attribute, default = None):
        value = unit.get_text_val(attribute, None)
        if value == None:
            baseunit = self.get_base_unit(unit)
            if baseunit:
                return self.get_unit_value(baseunit, attribute, default)
            return default
        return value

class UnitForest:
    """
    Contains the forest of unit advancement trees.
    """
    def __init__(self):
        self.trees = {}
        self.lookup = {}

    def add_node(self, un):
        """
        Add a new unit to the forest.
        """
        
        self.lookup[un.id] = un
        
        # First, we check if any of the new node's advancements already is in
        # the forest. If so, remove it and attach it to the new node.
        for cid in un.child_ids:
            if cid in self.trees:
                un.children.append(self.trees[cid])
                self.lookup[cid].parent_ids.append(un.id)
                del self.trees[cid]

        # Next, we check if the node is an advancement of an existing node. If
        # so, add it there.
        for rootnode in self.trees.values():
            if rootnode.try_add(un):
                return

        # Else, add a new tree with the new node as root.
        self.trees[un.id] = un

    def update_breadth(self):
        self.breadth = sum([x.update_breadth() for x in self.trees.values()])
        return self.breadth

    def get_children(self, uid):
        un = self.lookup[uid]
        return un.child_ids
    
    def get_parents(self, uid):
        un = self.lookup[uid]
        return un.parent_ids

class UnitNode:
    """
    A node in the advancement trees forest.
    """
    def __init__(self, unit):
        self.unit = unit
        self.children = []
        self.id = unit.get_text_val("id")
        self.child_ids = []
        self.parent_ids = []
        advanceto = unit.get_text_val("advanceto")
        if advanceto and advanceto != "null":
            for advance in advanceto.split(","):
                advance = advance.strip()
                self.child_ids.append(advance)

    def try_add(self, un):
        # A child of yours truly?
        if un.id in self.child_ids:
            self.children.append(un)
            un.parent_ids.append(self.id)
            return True
        # A recursive child?
        for child in self.children:
            if child.try_add(un): return True
        return False

    def update_breadth(self):
        if not self.children:
            self.breadth = 1
        else:
            self.breadth = sum([x.update_breadth() for x in self.children])
        return self.breadth   

class RaceNode:
    def __init__(self, race):
        self.race = race

html_header = '''
<!DOCTYPE HTML PUBLIC "-//W3C//DTD HTML 4.01//EN" "http://www.w3.org/TR/html4/strict.dtd">
<html>
<head>
<meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
<link rel=stylesheet href=\"%(path)sstyle.css\" type=\"text/css\">
</head>
<body>'''.strip()

html_footer = "</body></html>"

class HTMLOutput:
    def __init__(self, isocode, output, campaign, unitlist, verbose = False):
        self.isocode = isocode
        self.output = output
        self.campaign = campaign
        self.verbose = verbose
        self.target = "index.html"
        self.unitlist = unitlist
        self.forest = None

    def get_translation(self, textdomain, key):
        return self.unitlist.parser.translations.get(textdomain, self.isocode,
            key, key)

    def analyze_units(self): 
        # Build an advancement tree forest of all units.
        forest = self.forest = UnitForest()
        for u in self.unitlist.units_by_campaign[self.campaign]:
            forest.add_node(UnitNode(u))
        forest.update_breadth()

        # Partition trees by race of first unit.
        races = {}
        for tree in forest.trees.values():
            u = tree.unit
            # FIXME: use translated name for sorting
            race = u.get_text_val("race")
            if not race: continue
            races[race] = races.get(race, []) + [tree]        
        racelist = races.keys()
        racelist.sort()  

        rows_count = forest.breadth + len(racelist)
        # Create empty grid.
        rows = []
        for j in range(rows_count):
            column = []
            for i in range(6):
                column.append((1, 1, None))
            rows.append(column)

        # Sort advancement trees by name of first unit and place into the grid.
        def by_name(t1, t2):
            u1 = t1.unit
            u2 = t2.unit

            u1name = u1.get_text_val("name")
            u2name = u2.get_text_val("name")
            return cmp(u1name, u2name)

        def grid_place(nodes, x):
            nodes.sort(cmp = by_name)
            for node in nodes:
                level = int(node.unit.get_text_val("level"))
                rows[x][level] = (1, node.breadth, node)
                for i in range(1, node.breadth):
                    rows[x + i][level] = (0, 0, node)
                grid_place(node.children, x)
                x += node.breadth
            return x

        x = 0
        for race in racelist:
            node = RaceNode(self.unitlist.race_lookup[race])
            rows[x][0] = (6, 1, node)
            for i in range(1, 6):
                rows[x][i] = (0, 0, None)
            nodes = races[race]
            x += 1
            x = grid_place(nodes, x)
            
        self.racelist = racelist
        self.unitgrid = rows

    def write_navbar(self):
        def write(x): self.output.write(x)
        languages = find_languages()
        langlist = languages.keys()
        langlist.sort()
        write("<table class=\"navbar\">")

        write("<tr>\n")
        write("<th>By Campaign</th></tr><tr><td>")
        for campaign in all_campaigns:
            lang = self.isocode
            # FIXME: translate campaign names
            write(" <a href=\"../%s_%s/index.html\">%s</a><br/>\n" % (
                lang, campaign, campaign.replace("_", " ")))
        write("</td>\n")
        write("</tr>\n")

        write("<tr>\n")
        write("<th>By Faction</th></tr><tr><td>TODO</td>\n")
        write("</tr>\n")

        write("<tr>\n")
        write("<th>By Race</th></tr><tr><td>")
        write("<a href=\"index.html\">All</a><br/>\n")
        for rid in self.racelist:
            race = self.unitlist.race_lookup[rid]
            racename = race.get_text_val("plural_name")
            write(" <a href=\"index.html\">%s</a><br/>" % racename)
        write("</td>\n")
        write("</tr>\n")
        
        write("<tr>\n")
        write("<th>Language</th></tr><tr><td>")
        for lang in langlist:
            write(" <a href=\"../%s_%s/%s\">%s</a><br/>\n" % (
                lang, self.campaign, self.target, languages[lang]))
        write("</td>\n")
        write("</tr>\n")

        write("</table>\n")

    def pic(self, u, x):
        image = self.unitlist.get_unit_value(x, "image")  
        if not image:
            if x.name == "female":
                baseunit = self.unitlist.get_base_unit(u)
                if baseunit:
                    female = baseunit.get_first("female")
                    return self.pic(u, female)
            sys.stderr.write(
                "Warning: Missing image for unit %s(%s).\n" % (
                u.get_text_val("id"), x.name))
            return None
        image_collector.add(self.campaign, image)
        image = os.path.join("../pics/%s" % self.campaign,
            image)
        return image

    def write_units(self):
        def write(x): self.output.write(x)
        rows = self.unitgrid
        write("<table class=\"units\">\n")
        write("<colgroup>")
        for i in range(6):
            write("<col class=\"col%d\" />" % i)
        write("</colgroup>")
        for row in range(len(rows)):
            write("<tr>\n")
            for column in range(6):
                hspan, vspan, un = rows[row][column]
                if vspan:
                    attributes = ""
                    if hspan == 1 and vspan == 1:
                        pass
                    elif hspan == 1:
                        attributes += " rowspan=%d" % vspan
                    elif vspan == 1:
                        attributes += " colspan=%d" % hspan

                    if un and isinstance(un, RaceNode):
                        racename = un.race.get_text_val("plural_name")
                        attributes += " class=\"raceheader\""
                        write("<td%s>" % attributes)
                        write("%s" % racename)
                        write("</td>\n")
                    elif un:
                        u = un.unit
                        attributes += " class=\"unitcell\""
                        write("<td%s>" % attributes)

                        name = u.get_text_val("name")
                        
                        uid = u.get_text_val("id")
                        cost = u.get_text_val("cost")
                        hp = u.get_text_val("hitpoints")
                        mp = u.get_text_val("movement")
                        xp = u.get_text_val("experience")
                        level = u.get_text_val("level")
                        
                        write("<div class=\"l\">L%s</div>" % level)
                        link = "../%s_%s/%s.html" % (self.isocode,
                            self.campaign, uid)
                        write("<a href=\"%s\">%s</a><br/>" % (link, name))
                        
                        write('<div class="pic">')
                        image = self.pic(u, u)
                        write('<a href=\"%s\"><img src="%s" alt="(image)" /></a>\n' % (
                            link, image))
                        write('</div>\n')
                        write("<div class=\"attributes\">")
                        write("cost: %s<br />" % cost)
                        write("HP: %s<br />" % hp)
                        write("MP: %s<br />" % mp)
                        write("XP: %s<br />" % xp)
                        write("</div>")
                        write("</td>\n")
                    else:
                        write("<td class=\"empty\"></td>")
            write("</tr>\n")
        write("</table>\n")

    def write_unit_tree(self):
        self.output.write(html_header % {"path" : "../"})
       
        self.analyze_units()
        self.write_navbar()
        self.write_units()
            
        self.output.write(html_footer)

    def write_unit_report(self, output, unit):
        def write(x): self.output.write(x)

        self.output = output
        write(html_header % {"path" : "../"})
        self.write_navbar()

        # Write unit name, picture and description.
        uid = unit.get_text_val("id")
        uname = unit.get_text_val("name")
        display_name = uname

        female = unit.get_first("female")
        if female:
            fname = female.get_text_val("name")
            if fname != uname:
                display_name += "<br/>" + fname

        write("<h1>%s</h1>\n" % display_name)
        
        write('<div class="pic">')
        if female:
            mimage = self.pic(unit, unit)
            fimage = self.pic(unit, female)
            if not fimage: fimage = mimage
            write('<img src="%s" alt="(image)" />\n' % mimage)
            write('<img src="%s" alt="(image)" />\n' % fimage)
        else:
            image = self.pic(unit, unit)
            write('<img src="%s" alt="(image)" />\n' % image)
        write('</div>\n')
                        
        description = unit.get_text_val("description")
        if not description: description = "-"
        write("<p>%s</p>\n" % description)
        
        # Base info.
        
        hp = self.unitlist.get_unit_value(unit, "hitpoints")
        mp = self.unitlist.get_unit_value(unit, "movement")
        xp = self.unitlist.get_unit_value(unit, "experience")
        level = self.unitlist.get_unit_value(unit, "level")
        alignment = self.unitlist.get_unit_value(unit, "alignment")  

        write("<p>%s" % self.get_translation("wesnoth", "Advances from: "))
        for pid in self.forest.get_parents(uid):
            link = "../%s_%s/%s.html" % (self.isocode, self.campaign, pid)
            name = self.unitlist.unit_lookup[pid].get_text_val("name")
            write(" <a href=\"%s\">%s</a>" % (link, name))
        write("</p>\n")
        write("<p>%s" % self.get_translation("wesnoth", "Advances to: "))
        for cid in self.forest.get_children(uid):
            link = "../%s_%s/%s.html" % (self.isocode, self.campaign, cid)
            name = self.unitlist.unit_lookup[cid].get_text_val("name")
            write(" <a href=\"%s\">%s</a>" % (link, name))
        write("</p>\n")
        write("<table class=\"unitinfo\">\n")
        for val in ["cost", "hitpoints", "movement", "experience", "level",
            "alignment"]:
            write("<tr>\n")
            write("<th>%s</th>" % val)
            x = self.unitlist.get_unit_value(unit, val)
            write("<td>%s</td>" % x)
            write("</tr>\n")
        write("</table>\n")
        
        # Write info about movement costs and terrain defense.
        write("<table class=\"unitinfo\">\n")
        write("<tr>\n")
        write("<th>Terrain</th><th>Movement</th><th>Defense</th>\n")
        write("</tr>\n")
        
        def find_attr(what, key):
            if unit.movetype:
                mtx = unit.movetype.get_first(what)
                mty = None
                if mtx: mty = mtx.get_text_val(key)
            x = unit.get_first(what)
            y = None
            if x: y = x.get_text_val(key)
            if y:
                return True, y
            if unit.movetype and mty != None:
                return False, mty
            return False, "-"

        terrains = self.unitlist.terrain_lookup
        terrainlist = []
        for tid, t in terrains.items():
            if tid in ["off_map", "fog", "shroud"]: continue
            if t.get_first("aliasof"): continue
            name = t.get_text_val("name")
            terrainlist.append((name, tid))
        terrainlist.sort()

        for tname, tid in terrainlist:
            special, c = find_attr("movement_costs", tid)
            ccell = "td"
            if special: ccell += ' class="special"'
            dcell = "td"
            special, d = find_attr("defense", tid)
            if special: dcell += ' class="special"'

            write("<tr>\n")
            write("<td>%s</td><%s>%s</td><%s>%s</td>\n" % (
                tname, ccell, c, dcell, d))
            write("</tr>\n")
        write("</table>\n")
        
        # Write info about resistances.
        resistances = [
            "blade",
            "pierce",
            "impact",
            "fire",
            "cold",
            "arcane"]

        write("<table class=\"unitinfo\">\n")
        for rid in resistances:
            special, r = find_attr("resistance", rid)
            rcell = "td"
            if special: rcell += ' class="special"'
            write("<tr>\n")
            write("<th>%s</th><td>%s</td>\n" % (
                self.get_translation("wesnoth", rid), r))
            write("</tr>\n")
        write("</table>\n")
        
        # Write info about attacks.
        write("<table class=\"unitinfo\">\n")
        for attack in unit.get_all("attack"):
            write("<tr>")
            
            aid = attack.get_text_val("name")
            aname = attack.get_text_val("description")

            icon = attack.get_text_val("icon")
            if not icon:
                icon = "attacks/%s.png" % aid
            image_collector.add(self.campaign, icon)
            icon = os.path.join("../pics/%s" % self.campaign, icon)
            write("<td><img src=\"%s\" alt=\"(image)\"/></td>" % icon)

            write("<td>%s" % aname)
            
            t = attack.get_text_val("type")
            write("<br/>%s</td>" % self.get_translation("wesnoth", t))

            n = attack.get_text_val("number")
            x = attack.get_text_val("damage")
            try:
                x = " ".join(int(n) * [x])
            except TypeError:
                sys.stderr.write("Warning: Invalid damage for %s in %s.\n" % (
                    aname, uname.encode("utf8")))
                x = "?"
            write("<td>%s" % x)

            r = attack.get_text_val("range")
            write("<br/>%s</td>" % self.get_translation("wesnoth", r))
            
            s = []
            specials = attack.get_first("specials")
            if specials:
                for special in specials.children():
                    s.append(special.name)
            s = "<br/>".join(s)
            write("<td>%s</td>" % s)
            
            
            write("</tr>")
        write("</table>\n")

        write(html_footer)

languages_found = {}
def find_languages():
    """
    Returns a dictionary mapping isocodes to languages.
    """
    global languages
    if languages_found: return languages_found
    parser = wmlparser.Parser(datadir)
    WML = wmldata.DataSub("WML")
    parser.parse_text("{languages}\n")
    parser.parse_top(WML)

    for locale in WML.get_all("locale"):
        isocode = locale.get_text_val("locale")
        name = locale.get_text_val("name")
        languages_found[isocode] = name
    return languages_found

factions_found = {}
def find_factions():
    class Faction:
        def __init__(self, faction):
            self.faction = faction
    parser = wmlparser.Parser(datadir)
    WML = wmldata.DataSub("WML")
    parser.parse_text("{multiplayer/eras.cfg}\n")
    parser.parse_top(WML)
    
    for era in WML.get_all("era"):
        if era.get_text_val("id") == "era default":
            for multiplayer_side in era.get_all("multiplayer_side"):
                fid = multiplayer_side.get_text_val("id")
                factions_found[fid] = Faction(multiplayer_side)

class MyFile:
    """
    I don't understand why this is needed..
    """
    def __init__(self, filename, mode):
        self.f = open(filename, mode)
    def write(self, x):
        x = x.encode("utf8")
        self.f.write(x)

def generate_report(out_path, isocode, campaign, unitlist):
    if not campaign in unitlist.units_by_campaign: return
    print "Generating report for %s_%s." % (isocode, campaign)
    path = os.path.join(out_path, isocode + "_" + campaign)
    if not os.path.isdir(path): os.mkdir(path)
    output = MyFile(os.path.join(path, "index.html"), "w")
    html = HTMLOutput(isocode, output, campaign, unitlist)
    html.write_unit_tree()
    
    for unit in unitlist.units_by_campaign[campaign]:
        uid = unit.get_text_val("id")
        output = MyFile(os.path.join(path, "%s.html" % uid), "w")
        html.target = "%s.html" % uid
        html.write_unit_report(output, unit)        

def write_index(out_path):
    output = MyFile(os.path.join(out_path, "index.html"), "w")
    output.write(html_header % {"path" : ""})
    output.write("<div id=\"navbar\">\n")
    output.write("Units for:<ul>")
    for campaign in all_campaigns:
        output.write("<li><a href=\"C_%s/index.html\">%s</a></li>" % (
            campaign, campaign.replace("_", " ")))
    output.write("</ul>")
    output.write("</div>")
    output.write(html_footer)

def output(isocode):
    global all_campaigns

    unitlist = UnitList(isocode)

    # Parse all unit data
    # This reads in units.cfg, giving us all the mainline units.
    all_campaigns = []
    all_updated_campaigns = []
    
    print "Reading WML stuff."
    unitlist.add_terrains()

    print "Reading mainline units."
    unitlist.add("{core/units.cfg}", "mainline")
    all_updated_campaigns.append("mainline")
    all_campaigns.append("mainline")

    # Now we read each campaign in turn to get its units.
    campaigns = glob.glob("data/campaigns/*")
    for campaign in campaigns:
        dirname = campaign[5:] # strip leading data/
        description = dirname[10:]
        if options.campaign == "all" or options.campaign == description:
            print "Reading %s units." % description
            if unitlist.add("{%s}" % dirname, description):
                all_updated_campaigns.append(description)
                all_campaigns.append(description)

    # Report generation
    if not os.path.isdir(options.output):
        os.mkdir(options.output)
    write_index(options.output)

    if options.campaign == "all":
        campaigns = all_updated_campaigns
    else:
        campaigns = [options.campaign]

    if not campaigns:
        sys.stderr.write("No such campaign: %s\n" % options.campaign)
        sys.exit(1)    

    for campaign in campaigns:
        generate_report(options.output, isocode, campaign, unitlist)

if __name__ == '__main__':
    global options
    import optparse

    op = optparse.OptionParser()
    op.add_option("-c", "--campaign", default = "all",
        help = "Specify a campaign.")
    op.add_option("-l", "--language", default = "all",
        help = "Specify a language.")
    op.add_option("-o", "--output",
        help = "Specify output directory.")
    op.add_option("-n", "--nocopy", action = "store_true",
        help = "No copying of files.")
    options, args = op.parse_args()
    
    if not options.output:
        op.print_help()
        sys.exit(-1)

    wmltools.pop_to_top("wmlunits")
    datadir = os.getcwd() + "/data"
    
    if options.language == "all":
        languages = find_languages().keys()
    else:
        languages = [options.language]
    
    if not languages:
        sys.stderr.write("No such language: %s\n" % options.language)
        sys.exit(1)  
    
    for isocode in languages:
        output(isocode)

    if not options.nocopy:
        print "Copying files."
        image_collector.copy_and_color_images(options.output)
        shutil.copy2("data/tools/unit_tree/style.css", options.output)