Permalink
Browse files

Make the new unit back-compatible via delegation.

  • Loading branch information...
1 parent 22a5040 commit 1954dfc36fa847902cd93a3534e05eb38d49eac4 @GraylinKim committed Apr 3, 2013
Showing with 45 additions and 11 deletions.
  1. +45 −11 sc2reader/data/__init__.py
View
@@ -27,21 +27,56 @@ class Unit(object):
def __init__(self, unit_id, flags):
self.id = unit_id
self.flags = flags
- self.type = None
- self.name = 'Unknown Unit'
+ self._type_class = None
self.hallucinated = (flags & 2 == 2)
def is_type(self, unit_type):
if isinstance(unit_type, int):
- if self.type:
- return self.type.id == unit_type
- else:
- return unit_type == 0
+ # Compare integer ids. Unknown units have id==0 and should be equal
+ return self._type_class.id if self._type_class else unit_type == 0
else:
- return self.type == unit_type
+ return self._type_class == unit_type
- def get_race(self):
- return self.type.race if self.type else None
+ @property
+ def name(self):
+ return self._type_class.name if self._type_class else None
+
+ @property
+ def title(self):
+ return self._type_class.title if self._type_class else None
+
+ @property
+ def type(self):
+ """ For backwards compatibility this returns the int id instead of the actual class """
+ return self._type_class.id if self._type_class else None
+
+ @property
+ def race(self):
+ return self._type_class.race if self._type_class else None
+
+ @property
+ def minerals(self):
+ return self._type_class.minerals if self._type_class else None
+
+ @property
+ def vespene(self):
+ return self._type_class.vespene if self._type_class else None
+
+ @property
+ def supply(self):
+ return self._type_class.supply if self._type_class else None
+
+ @property
+ def is_worker(self):
+ return self._type_class.is_worker if self._type_class else False
+
+ @property
+ def is_building(self):
+ return self._type_class.is_building if self._type_class else False
+
+ @property
+ def is_army(self):
+ return self._type_class.is_army if self._type_class else False
def __str__(self):
return "{0} [{1:X}]".format(self.name, self.id)
@@ -74,8 +109,7 @@ def create_unit(self, unit_id, unit_type, unit_flags):
def change_type(self, unit, new_type):
if new_type in self.units:
reference_unit = self.units[new_type]
- unit.type = reference_unit
- unit.name = reference_unit.name
+ unit._type_class = reference_unit
else:
self.logger.error("Unable to change type of {0} to {1}; unit type not found in build {2}".format(unit,hex(new_type),self.id))

0 comments on commit 1954dfc

Please sign in to comment.