home *** CD-ROM | disk | FTP | other *** search
/ Maximum CD 2010 July / maximum-cd-2010-07.iso / DiscContents / wesnoth-1.8-win32.exe / data / tools / unit_tree / helpers.py < prev    next >
Encoding:
Python Source  |  2010-01-17  |  16.3 KB  |  479 lines
  1. """
  2. Various helpers for use by the wmlunits tool.
  3. """
  4. import sys, os, re, glob, shutil, copy, urllib2
  5. from subprocess import Popen
  6.  
  7. import wesnoth.wmldata as wmldata
  8. import wesnoth.wmlparser as wmlparser
  9. import wesnoth.wmltools as wmltools
  10.  
  11. class ParserWithCoreMacros:
  12.     """
  13.     A wrapper around the WML parser to do some things like we want.
  14.     """
  15.     def __init__(self, isocode, datadir, userdir, transdir):
  16.         self.datadir = datadir
  17.         self.userdir = userdir
  18.         # Handle translations.
  19.         self.translations = wmltools.Translations(transdir)
  20.         def gettext(textdomain, x):
  21.             return self.translations.get(textdomain, isocode, x, x)
  22.         self.gettext = gettext
  23.  
  24.         # Create a new parser for the macros.
  25.         parser = wmlparser.Parser(datadir)
  26.         parser.gettext = self.gettext
  27.         
  28.         # Parse core macros.
  29.         parser.parse_text("{core/macros/}\n")
  30.         parser.parse_top(None)
  31.         self.core_macros = parser.macros
  32.                 
  33.     def parse(self, text_to_parse, ignore_macros = None,
  34.         ignore_fatal_errors = False, verbose = False):
  35.  
  36.         # Create the real parser.
  37.         parser = wmlparser.Parser(self.datadir, self.userdir)
  38.         parser.verbose = verbose
  39.         parser.gettext = self.gettext
  40.         parser.macros = copy.copy(self.core_macros)
  41.         
  42.         #parser.verbose = True
  43.         
  44.         # Suppress complaints about undefined terrain macros
  45.         parser.set_macro_not_found_callback(lambda wmlparser, name, params:
  46.         name.startswith("TERRAIN") or name == "DISABLE_TRANSITIONS")
  47.         
  48.         if ignore_macros:
  49.             parser.macro_callback = ignore_macros
  50.  
  51.         # Create a WML root element and parse the given text into it.
  52.         WML = wmldata.DataSub("WML")
  53.  
  54.         parser.parse_text(text_to_parse)
  55.  
  56.         parser.parse_top(WML)
  57.  
  58.         return WML
  59.  
  60. class ImageCollector:
  61.     """
  62.     A class to collect all the images which need to be copied to the HTML
  63.     output folder.
  64.     """
  65.     def __init__(self, datadir, userdir):
  66.         self.images = {}
  67.         self.pathes_per_campaign = {}
  68.         self.ipathes = {}
  69.         self.notfound = {}
  70.         self.datadir = datadir
  71.         self.userdir = userdir
  72.         self.id = 0
  73.         self.verbose = 0
  74.  
  75.     def add_binary_pathes_from_WML(self, campaign, WML):
  76.         self.pathes_per_campaign[campaign] = self.pathes_per_campaign.get(
  77.             campaign, [])
  78.         for binpath in WML.get_all("binary_path"):
  79.             path = binpath.get_text_val("path")
  80.             self.pathes_per_campaign[campaign].append(path)
  81.  
  82.     def find_image(self, i, c):
  83.         if c == "mainline":
  84.             bases = [os.path.join(self.datadir, "core/images")]
  85.         else:
  86.             bases = [os.path.join(self.datadir, "core/images")]
  87.             binpaths = self.pathes_per_campaign.get(c, [])
  88.             binpaths.reverse()
  89.             for x in binpaths:
  90.                 if x.startswith("data/"):
  91.                     for idir in ["images", "images/units"]:
  92.                         bases.append(os.path.join(self.datadir, x[5:], idir))
  93.                         if self.userdir:
  94.                             bases.append(os.path.join(self.userdir, x[5:], idir))
  95.         for base in bases:
  96.             ipath = os.path.join("%s" % base, i)
  97.             if os.path.exists(ipath): return ipath, None
  98.         return None, bases
  99.  
  100.     def add_image_check(self, campaign, path):
  101.         if (campaign, path) in self.notfound:
  102.             return self.notfound[(campaign, path)], True
  103.         ipath, error = self.find_image(path, campaign)
  104.         if ipath in self.ipathes:
  105.             return self.ipathes[ipath], False
  106.     
  107.         name = "%05d_" % self.id
  108.         name += os.path.basename(path)
  109.         self.id += 1
  110.  
  111.         self.images[name] = ipath, path, campaign, error
  112.         if ipath:
  113.             self.ipathes[ipath] = name
  114.             return name, False
  115.         else:
  116.             self.notfound[(campaign, path)] = name
  117.             return name, True
  118.         
  119.     def add_image(self, campaign, path):
  120.         name, error = self.add_image_check(campaign, path)
  121.         return name
  122.  
  123.     def copy_and_color_images(self, target_path):
  124.         for iid in self.images.keys():
  125.             opath = os.path.join(target_path, "pics", iid)
  126.             try:
  127.                 os.makedirs(os.path.dirname(opath))
  128.             except OSError:
  129.                 pass
  130.  
  131.             ipath, i, c, bases = self.images[iid]
  132.             if ipath and os.path.exists(ipath):
  133.                 #shutil.copy2(ipath, opath)
  134.                 # We assume TeamColorizer is in the same directory as the
  135.                 # helpers.py currently executing.
  136.                 command = os.path.join(os.path.dirname(__file__),
  137.                     "TeamColorizer")
  138.                 p = Popen([command, ipath, opath])
  139.                 p.wait()
  140.             else:
  141.                 sys.stderr.write(
  142.                     "Warning: Required image %s: \"%s\" does not exist.\n" % (
  143.                         repr(c), repr(i)))
  144.                 if self.verbose:
  145.                     sys.stderr.write("Warning: Looked at the following locations:\n")
  146.                     sys.stderr.write("\n".join(bases) + "\n")
  147.  
  148. blah = 1
  149. class WesnothList:
  150.     """
  151.     Lists various Wesnoth stuff like units, campaigns, terrains, factions...
  152.     """
  153.     def __init__(self, isocode, datadir, userdir, transdir):
  154.         self.unit_lookup = {}
  155.         self.race_lookup = {}
  156.         self.terrain_lookup = {}
  157.         self.movetype_lookup = {}
  158.         self.era_lookup = {}
  159.         self.campaign_lookup = {}
  160.         self.parser = ParserWithCoreMacros(isocode, datadir, userdir, transdir)
  161.  
  162.     def add_terrains(self):
  163.         """
  164.         We need those for movement costs and defenses.
  165.         """
  166.         WML = self.parser.parse("{core/terrain.cfg}\n")
  167.  
  168.         n = 0
  169.         for terrain in WML.get_all("terrain_type"):
  170.             tstring = terrain.get_text_val("string")
  171.             self.terrain_lookup[tstring] = terrain
  172.             n += 1
  173.         return n
  174.  
  175.     def add_era(self, era):
  176.         """
  177.         For an era, list all factions and units in it.
  178.         """
  179.         eid = era.get_text_val("id")
  180.         if not eid: return
  181.         self.era_lookup[eid] = era
  182.         era.faction_lookup = {}
  183.         for multiplayer_side in era.get_all("multiplayer_side"):
  184.             fid = multiplayer_side.get_text_val("id")
  185.             if fid == "Random": continue
  186.             era.faction_lookup[fid] = multiplayer_side
  187.             recruit = multiplayer_side.get_text_val("recruit", "").strip()
  188.             leader = multiplayer_side.get_text_val("leader", "").strip()
  189.             units = recruit.split(",")
  190.             leaders = leader.split(",")
  191.             multiplayer_side.units = {}
  192.             multiplayer_side.is_leader = {}
  193.             for u in units:
  194.                 uid = u.strip()
  195.                 if uid:
  196.                     multiplayer_side.units[uid] = True
  197.             for u in leaders:
  198.                 uid = u.strip()
  199.                 if uid:
  200.                     multiplayer_side.units[uid] = True
  201.                     multiplayer_side.is_leader[uid] = True
  202.         return eid
  203.  
  204.     def add_campaign(self, campaign):
  205.         name = campaign.get_text_val("id")
  206.         if not name:
  207.             global blah
  208.             name = "noid%d" % blah
  209.             blah += 1
  210.         self.campaign_lookup[name] = campaign
  211.         return name
  212.  
  213.     def add_mainline_eras(self):
  214.         """
  215.         Find all mainline eras.
  216.         """
  217.         WML = self.parser.parse("{multiplayer/eras.cfg}\n")
  218.  
  219.         n = 0
  220.         for era in WML.get_all("era"):
  221.             self.add_era(era)
  222.             n += 1
  223.         return n
  224.  
  225.     def add_mainline_campaigns(self):
  226.         """
  227.         Find all mainline campaigns.
  228.         """
  229.         WML = self.parser.parse("{campaigns}\n")
  230.         n = 0
  231.         for campaign in WML.find_all("campaign"):
  232.             self.add_campaign(campaign)
  233.             n += 1
  234.         return n
  235.  
  236.     def add_addons(self, image_collector):
  237.         """
  238.         Find all addon eras and campaigns.
  239.         """
  240.         n = 0
  241.         try:
  242.             WML = self.parser.parse("""
  243.                 #define MULTIPLAYER\n#enddef
  244.                 #define RANDOM_SIDE\n#enddef
  245.                 {~add-ons}
  246.                 """)
  247.         except wmlparser.Error, e:
  248.             print e
  249.             return n
  250.         for campaign in WML.find_all("campaign"):
  251.             cid = self.add_campaign(campaign)
  252.  
  253.         for era in WML.find_all("era"):
  254.             eid = self.add_era(era)
  255.             if not eid: continue
  256.             image_collector.add_binary_pathes_from_WML(eid, WML)
  257.             
  258.         n = self.add_units(WML, "addons")
  259.         
  260.         image_collector.add_binary_pathes_from_WML("addons", WML)
  261.         return n
  262.  
  263.     def add_units(self, WML, campaign):
  264.         """
  265.         We assume each unit, in mainline and all addons, has one unique id. So
  266.         we reference them everywhere by this id, and here can add them all to
  267.         one big collection.
  268.         """
  269.         addunits = WML.get_all("units")
  270.         if not addunits: return 0
  271.         
  272.         def getall(oftype):
  273.             r = []
  274.             for units in addunits:
  275.                 r += units.get_all(oftype)
  276.             return r
  277.  
  278.         # Find all unit types.
  279.         newunits = getall("unit_type") + getall("unit")
  280.         for unit in newunits:
  281.             if unit.get_text_val("do_not_list", "no") == "no" and\
  282.                unit.get_text_val("hide_help", "no") in ["no", "false"]:
  283.                 uid = unit.get_text_val("id")
  284.                 if uid in self.unit_lookup:
  285.                     sys.stderr.write(
  286.                         ("Fatal: Unit id \"%s\" already exists - either it has" +
  287.                         " to be renamed, or there's a bug in this script.\n") % 
  288.                         uid)
  289.                 else:
  290.                     self.unit_lookup[uid] = unit
  291.                 unit.campaign = campaign
  292.  
  293.         # Find all races.
  294.         newraces = getall("race")
  295.         for race in newraces:
  296.             rid = race.get_text_val("id")
  297.             if rid == None:
  298.                 rid = race.get_text_val("name")
  299.  
  300.             self.race_lookup[rid] = race
  301.  
  302.         # Find all movetypes.
  303.         newmovetypes = getall("movetype")
  304.         for movetype in newmovetypes:
  305.             mtname = movetype.get_text_val("name")
  306.             self.movetype_lookup[mtname] = movetype
  307.  
  308.         # Store race/movetype/faction of each unit for easier access later.
  309.         for unit in newunits:
  310.             uid = unit.get_text_val("id")
  311.             race = self.get_unit_value(unit, "race")
  312.             try: unit.race = self.race_lookup[race]
  313.             except KeyError:
  314.                 unit.race = None
  315.                 sys.stderr.write("Warning: No race \"%s\" found (%s).\n" % (
  316.                     race, unit.get_text_val("id")))
  317.             movetype = self.get_unit_value(unit, "movement_type")
  318.             try: unit.movetype = self.movetype_lookup[movetype]
  319.             except KeyError: unit.movetype = None
  320.             
  321.             unit.advance = []
  322.             advanceto = unit.get_text_val("advances_to")
  323.             # Add backwards compatibility for 1.4
  324.             if not advanceto:
  325.                 advanceto = unit.get_text_val("advanceto")
  326.             if advanceto and advanceto != "null":
  327.                 for advance in advanceto.split(","):
  328.                     auid = advance.strip()
  329.                     if auid: unit.advance.append(auid)
  330.  
  331.             # level
  332.             try:
  333.                 level = int(self.get_unit_value(unit, "level"))
  334.             except TypeError:
  335.                 level = 0
  336.             except ValueError:
  337.                 level = 0
  338.             if level < 0: level = 0
  339.             unit.level = level
  340.  
  341.         return len(newunits)
  342.  
  343.     def find_unit_factions(self):
  344.         for unit in self.unit_lookup.values():
  345.             unit.factions = []
  346.             unit.eras = []
  347.  
  348.         for eid, era in self.era_lookup.items():
  349.             for fid, multiplayer_side in era.faction_lookup.items():
  350.                 for uid in multiplayer_side.units:
  351.                     try:
  352.                         unit = self.unit_lookup[uid]
  353.                     except KeyError:
  354.                         sys.stderr.write(
  355.                             ("Error: Era '%s' faction '%s' references " +
  356.                             "non-existant unit id '%s'!\n") % (
  357.                                 eid,
  358.                                 fid,
  359.                                 repr(uid)))
  360.                         continue
  361.                     if not eid in unit.eras:
  362.                         unit.eras.append(eid)
  363.                     unit.factions.append((eid, fid))
  364.  
  365.     def get_base_unit(self, unit):
  366.         b = unit.get_first("base_unit")
  367.         if b:
  368.             buid = b.get_text_val("id")
  369.             try: baseunit = self.unit_lookup[buid]
  370.             except KeyError:
  371.                 sys.stderr.write(
  372.                     "Warning: No baseunit \"%s\" for \"%s\".\n" % (
  373.                     buid, unit.get_text_val("id")))
  374.                 return None
  375.             return baseunit
  376.         return None
  377.  
  378.     def get_unit_value(self, unit, attribute, default = None):
  379.         value = unit.get_text_val(attribute, None)
  380.         if value == None:
  381.             baseunit = self.get_base_unit(unit)
  382.             if baseunit:
  383.                 return self.get_unit_value(baseunit, attribute, default)
  384.             return default
  385.         return value
  386.  
  387. class UnitForest:
  388.     """
  389.     Contains the forest of unit advancement trees.
  390.     """
  391.     def __init__(self):
  392.         self.trees = {}
  393.         self.lookup = {}
  394.  
  395.     def add_node(self, un):
  396.         """
  397.         Add a new unit to the forest.
  398.         """
  399.         self.lookup[un.id] = un
  400.  
  401.     def create_network(self):
  402.         """
  403.         Assuming that each node which has been added to the tree only has a
  404.         valid list of children in unit.child_ids, also fill in unit.parent_ids
  405.         and update the unit.children shortcut.
  406.         """
  407.         
  408.         # Complete the network
  409.         for uid, u in self.lookup.items():
  410.             for cid in u.child_ids:
  411.                 c = self.lookup.get(cid, None)
  412.                 if not c: continue
  413.                 u.children.append(c)
  414.                 if not uid in c.parent_ids:
  415.                     c.parent_ids.append(uid)
  416.             
  417.  
  418.         # Put all roots into the forest
  419.         for uid, u in self.lookup.items():
  420.             if not u.parent_ids:
  421.                 self.trees[uid] = u
  422.         
  423.         # Sanity check because some argGRRxxx addons have units who advance to
  424.         # themselves.
  425.  
  426.         def recurse(u, already):
  427.             already2 = already.copy()
  428.             for c in u.children[:]:
  429.                 already2[c.id] = True
  430.                 if c.id in already:
  431.                     sys.stderr.write(
  432.                         "Warning: Unit %s advances to unit %s in a loop.\n" %
  433.                         (u.id, c.id))
  434.                     sys.stderr.write("    Removing advancement %s.\n" % c.id)
  435.                     u.children.remove(c)
  436.             for c in u.children:
  437.                 recurse(c, already2)
  438.         for u in self.trees.values():
  439.             already = {u.id : True}
  440.             recurse(u, already)
  441.  
  442.     def update(self):
  443.         self.create_network()       
  444.     
  445.         self.breadth = sum([x.update_breadth() for x in self.trees.values()])
  446.         return self.breadth
  447.  
  448.     def get_children(self, uid):
  449.         un = self.lookup[uid]
  450.         return un.child_ids
  451.     
  452.     def get_parents(self, uid):
  453.         un = self.lookup[uid]
  454.         return un.parent_ids
  455.  
  456. class UnitNode:
  457.     """
  458.     A node in the advancement trees forest.
  459.     """
  460.     def __init__(self, unit):
  461.         self.unit = unit
  462.         self.children = []
  463.         self.id = unit.get_text_val("id")
  464.         self.child_ids = []
  465.         self.parent_ids = []
  466.         self.child_ids.extend(unit.advance)
  467.  
  468.     def update_breadth(self):
  469.         if not self.children:
  470.             self.breadth = 1
  471.         else:
  472.             self.breadth = sum([x.update_breadth() for x in self.children])
  473.         return self.breadth   
  474.  
  475. class GroupNode:
  476.     def __init__(self, data):
  477.         self.data = data
  478.  
  479.