keakon的涂鸦馆

# -*- coding: utf-8 -*-

try:
  from lxml import etree
except ImportError:
  try:
    import xml.etree.cElementTree as etree
  except ImportError:
    import xml.etree.ElementTree as etree

from copy import deepcopy


class Parser(object):
  elements = []
  metas = []
  meta_arrays = []
  all_metas = []
  table_heade = None
  template = {}

  @classmethod
  def meta_element_handler(cls, meta_element, obj):
    value = meta_element.get('value')
    tag = meta_element.tag
    if tag == 'EditorCategories':
      value = value.split(',')[0].split(':')[1]
    obj[tag] = value

  @classmethod
  def meta_array_handler(cls, meta_array_element, meta_array, obj):
    obj[meta_array_element.get('index')] = meta_array_element.get('value')

  @classmethod
  def other_handler(cls, element, obj):
    pass

  @classmethod
  def filter(cls, obj, element):
    return True

  @classmethod
  def parse(cls):
    objs = {}
    template = cls.template
    metas = cls.metas
    meta_arrays = cls.meta_arrays

    for element in cls.elements:
      id = element.get('id')
      if not id:
        continue

      parent = element.get('parent')
      if parent:
        try:
          obj = deepcopy(objs[parent])
        except:
          obj = deepcopy(template)
      else:
        obj = deepcopy(template)

      if not cls.filter(obj, element):
        continue

      for meta in metas:
        meta_element = element.find(meta)
        if meta_element is not None:
          cls.meta_element_handler(meta_element, obj)

      for meta_array in meta_arrays:
        for meta_array_element in element.findall(meta_array):
          cls.meta_array_handler(meta_array_element, meta_array, obj)

      cls.other_handler(element, obj)

      objs[id] = obj

    return objs

  @classmethod
  def process_additional_data(cls, data, line):
    pass

  @classmethod
  def parse_to_file(cls, name, append_header='', quote_metas=()):
    objs = cls.parse()

    lines = ['%s,%s%s' % (name, ','.join(cls.table_header), append_header)]

    for id, data in objs.iteritems():
      line = [id]
      for meta in cls.all_metas:
        if meta in quote_metas:
          line.append('"%s"' % data[meta])
        else:
          line.append(data[meta])

      cls.process_additional_data(data, line)

      lines.append(','.join(line))

    csvfile = open(name + '.csv', 'w')
    csvfile.write('\n'.join(lines))
    csvfile.close()


class UnitParser(Parser):
  elements = etree.parse('UnitData.xml').getroot().findall('CUnit')

  metas = [
    'Race',
    'EditorCategories',
    'LifeMax',
    'LifeRegenRate',
    'ShieldsMax',
    'ShieldRegenRate',
    'ShieldRegenDelay',
    'EnergyStart',
    'EnergyMax',
    'EnergyRegenRate',
    'Sight',
    'AttackTargetPriority',
    'Radius',
    'InnerRadius',
    'SeparationRadius',
    'MinimapRadius',
    'SubgroupPriority',
    'LifeArmor',
    'Speed',
    'Acceleration',
    'LateralAcceleration',
    'StationaryTurningRate',
    'TurningRate',
    'Mover',
    'Height',
    'VisionHeight',
    'RepairTime',
    'Food',
    'CargoSize',
    'Mass',
  ]

  meta_arrays = [
    'CostResource',
    'PlaneArray',
    'Collide',
    'FlagArray',
    'Attributes',
  ]

  no_prefix_unit_meta_arrays = meta_arrays[-2:]

  unit_metas_in_arrays = {}
  for unit_element in elements:
    for unit_meta_array in meta_arrays:
      for unit_meta_array_element in unit_element.findall(unit_meta_array):
        index = unit_meta_array_element.get('index')
        if not (unit_meta_array == 'FlagArray' and index.startswith('AI')):
          unit_metas_in_arrays[(unit_meta_array + ' - ' + index) if unit_meta_array not in no_prefix_unit_meta_arrays else index] = 1

  def cmp_unit_meta_arrays(x, y):
    def priority(s):
      if s.startswith('CostResource'):
        return 0
      elif s.startswith('Collide'):
        return 1
      elif s.startswith('PlaneArray'):
        return 2
      return 3
    return cmp(priority(x), priority(y)) or cmp(x, y)

  additional_unit_metas = sorted(unit_metas_in_arrays.keys(), cmp=cmp_unit_meta_arrays)
  all_metas = metas + additional_unit_metas
  table_header = all_metas[:]
  table_header[1] = 'ObjectType'

  template = dict(zip(all_metas, [''] * len(all_metas)))

  @classmethod
  def meta_array_handler(cls, meta_array_element, meta_array, obj):
    index = meta_array_element.get('index')
    obj[(meta_array + ' - ' + index) if meta_array not in cls.no_prefix_unit_meta_arrays else index] = meta_array_element.get('value')

  @classmethod
  def other_handler(cls, element, obj):
    weapons = []
    for weapon in element.findall('WeaponArray[@Link]'):
      weapons.append(weapon.get('Link'))
    obj['WeaponArray'] = weapons

  @classmethod
  def process_additional_data(cls, data, line):
    line.append('"%s"' % ','.join(data['WeaponArray']))


UnitParser.parse_to_file('Unit', ',WeaponArray')


class WeaponParser(Parser):
  elements = etree.parse('WeaponData.xml').getroot().findall('CWeaponLegacy')

  metas = [
    'EditorCategories',
    'TargetFilters',
    'Period',
    'RandomDelayMin',
    'RandomDelayMax',
    'DamagePoint',
    'AllowedMovement',
    'Backswing', # 击退？
    'Arc',
    'ArcSlop',
    'Range',
    'MinimumRange',
    'MinScanRange',
    'RangeSlop',
    'DisplayAttackCount',
  ]
  meta_arrays = [
    'Options',
    'LegacyOptions',
  ]

  weapon_metas_in_arrays = {}
  for weapon_element in elements:
    for weapon_meta_array in meta_arrays:
      for weapon_meta_array_element in weapon_element.findall(weapon_meta_array):
        weapon_metas_in_arrays[weapon_meta_array_element.get('index')] = 1

  additional_weapon_metas = weapon_metas_in_arrays.keys()
  all_metas = metas + additional_weapon_metas
  table_header = all_metas[:]
  table_header[0] = 'Race'

  template = dict(zip(all_metas, [''] * len(all_metas)))


WeaponParser.parse_to_file('Weapon', quote_metas=['TargetFilters'])


class DamageParser(Parser):
  all_metas = metas = [
    'EditorCategories',
    'Amount',
    'ArmorReduction',
    'Kind',
    'Visibility',
    'SearchFilters',
  ]

  meta_arrays = [
    'AreaArray',
    'ExcludeArray',
    #'AttributeBonus',
  ]

  damage_root = etree.parse('EffectData.xml').getroot()
  elements = damage_root.findall('CEffectDamage')

  damage_bonus = {}
  for damage_bonus_element in damage_root.findall('.//AttributeBonus'):
    damage_bonus[damage_bonus_element.get('index')] = 1
  damage_bonus = ['AttributeBonus - ' + key for key in damage_bonus.keys()]

  template = dict(zip(metas, [''] * len(metas)) + [(meta_array, []) for meta_array in meta_arrays[:2]] + zip(damage_bonus, [''] * len(damage_bonus)))

  table_header = metas + damage_bonus + ['Radius', 'Fraction'] * 3 + ['ExcludeArray']
  table_header[0] = 'Race'

  @classmethod
  def meta_array_handler(cls, meta_array_element, meta_array, obj):
    if meta_array_element.get('removed') == "1":
      del obj[meta_array_element.tag][int(meta_array_element.get('index'))]
    elif meta_array == 'AreaArray':
      obj['AreaArray'].append((meta_array_element.get('Radius'), meta_array_element.get('Fraction')))
    elif meta_array == 'ExcludeArray':
      obj['ExcludeArray'].append(meta_array_element.get('Value'))

  @classmethod
  def other_handler(cls, element, obj):
    for bonus in element.findall('AttributeBonus'):
      obj['AttributeBonus - ' + bonus.get('index')] = bonus.get('value')

  @classmethod
  def filter(cls, obj, element):
    return obj['Amount'] or element.find('Amount') is not None

  @classmethod
  def process_additional_data(cls, data, line):
    for damage_bonus_item in cls.damage_bonus:
      line.append(data.get(damage_bonus_item, ''))

    area_array = data['AreaArray']
    for (radius, fraction) in area_array:
      line.append(radius)
      line.append(fraction)

    for i in xrange(len(area_array) - 3):
      line.append('')
      line.append('')

    line.append('"%s"' % ','.join(data['ExcludeArray']))


DamageParser.parse_to_file('Damage', quote_metas=['SearchFilters'])


class UpgradeParser(Parser):
  elements = etree.parse('UpgradeData.xml').getroot().findall('CUpgrade')

  all_metas = metas = [
    'Race',
  ]

  meta_arrays = [
    'EffectArray',
  ]

  template = {'Race': '', 'EffectArray': []}

  table_header = metas + ['Effect - ' + `i` for i in xrange(1, 65)]

  @classmethod
  def meta_array_handler(cls, meta_array_element, meta_array, obj):
    if not meta_array_element.get('Operation'):
      obj['EffectArray'].append('"%s: %s"' % (meta_array_element.get('Reference'), meta_array_element.get('Value')))

  @classmethod
  def filter(cls, obj, element):
    return obj['EffectArray'] or element.find('EffectArray') is not None

  @classmethod
  def process_additional_data(cls, data, line):
    line.extend(data['EffectArray'])


UpgradeParser.parse_to_file('Upgrade')


class BehaviorParser(Parser):
  root = etree.parse('BehaviorData.xml').getroot()
  elements = root.findall('CBehaviorBuff')

  metas = [
    'EditorCategories',
    'Duration',
  ]

  meta_arrays = []

  modifications = set()
  for modification in root.findall('.//Modification'):
    modifications |= set(modification.keys())
  modifications = list(modifications)

  vital_regen_arrays  = set()
  for vital_regen_array in root.findall('CBehaviorBuff/Modification/VitalRegenArray'):
    vital_regen_arrays.add('VitalRegenArray - ' + vital_regen_array.get('index'))
  vital_regen_arrays = list(vital_regen_arrays)

  all_metas = metas + modifications + vital_regen_arrays

  template = dict(zip(all_metas, [''] * len(all_metas)))

  table_header = all_metas[:]
  table_header[0] = 'Race'

  @classmethod
  def filter(cls, obj, element):
    return element.find('Modification') is not None or obj['Duration'] or element.find('Duration') is not None

  @classmethod
  def other_handler(cls, element, obj):
    modification = element.find('Modification')
    if modification is not None:
      obj.update(modification.attrib)

    for vital_regen_array in element.findall('Modification/VitalRegenArray'):
      obj['%s - %s' % (vital_regen_array.tag, vital_regen_array.get('index'))] = vital_regen_array.get('value')


BehaviorParser.parse_to_file('Behavior', quote_metas=['RadarFilters'])


abil_root = etree.parse('AbilData.xml').getroot()

class AbilParser(Parser):
  elements = []

  table_header = all_metas = []

  template = dict(zip(all_metas, [''] * len(all_metas)))

  @classmethod
  def parse(cls):
    objs = {}
    template = cls.template

    for element in cls.elements:

      parent = element.get('parent')
      if parent:
        try:
          obj = objs[parent].copy()
        except:
          obj = template.copy()
      else:
        obj = template.copy()

      editor_categories_element = element.find('EditorCategories')
      if editor_categories_element is not None:
        obj['Race'] = editor_categories_element.get('value').split(',')[0].split(':')[1]

      cls.other_handler(element, obj, objs)

    return objs


class TrainParser(AbilParser):
  elements = abil_root.findall('CAbilTrain')

  table_header = all_metas = ['Race', 'Time', 'Minerals', 'Vespene']

  template = dict(zip(all_metas, [''] * len(all_metas)))

  @classmethod
  def other_handler(cls, element, obj, objs):
    for info_array in element.findall('InfoArray'):
      id = info_array.find('Unit')
      if id is not None:
        temp_obj = obj.copy()
        temp_obj['Time'] = info_array.get('Time')
        for resource in info_array.findall('Resource'):
          temp_obj[resource.get('index')] = resource.get('value')

        objs[id.get('value')] = temp_obj


TrainParser.parse_to_file('Train')


class WarpTrainParser(AbilParser):
  elements = abil_root.findall('CAbilWarpTrain')

  table_header = all_metas = ['Race', 'Cooldown', 'Time', 'Minerals', 'Vespene']

  template = dict(zip(all_metas, [''] * len(all_metas)))

  @classmethod
  def other_handler(cls, element, obj, objs):
    for info_array in element.findall('InfoArray'):
      id = info_array.get('Unit')
      if id:
        temp_obj = obj.copy()
        temp_obj['Time'] = info_array.get('Time')
        for resource in info_array.findall('Resource'):
          temp_obj[resource.get('index')] = resource.get('value')
        temp_obj['Cooldown'] = info_array.find('Cooldown').get('TimeUse')

        objs[id] = temp_obj


WarpTrainParser.parse_to_file('WarpTrain')


class ResearchParser(AbilParser):
  elements = abil_root.findall('CAbilResearch')

  table_header = all_metas = ['Race', 'Time', 'Minerals', 'Vespene']

  template = dict(zip(all_metas, [''] * len(all_metas)))

  @classmethod
  def other_handler(cls, element, obj, objs):
    for info_array in element.findall('InfoArray'):
      id = info_array.get('Upgrade')
      if id:
        temp_obj = obj.copy()
        temp_obj['Time'] = info_array.get('Time')
        for resource in info_array.findall('Resource'):
          temp_obj[resource.get('index')] = resource.get('value')

        objs[id] = temp_obj


ResearchParser.parse_to_file('Research')