Skip to content

Commit cba6c3b

Browse files
GraylinKimGraylinKim
committed
Identifies build abilities and maps to the unit.
Adds: AbilityEvent.is_build = True|False Abilityevent.build_unit = UnitClass
1 parent ab1b0f4 commit cba6c3b

File tree

6 files changed

+228
-60
lines changed

6 files changed

+228
-60
lines changed

‎sc2reader/data/__init__.py

Lines changed: 204 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,13 @@
77
#from sc2reader.data.build19595 import Data_19595
88

99

10+
class Build(object):
11+
def __init__(self, build_id, units, abilities):
12+
self.id=build_id
13+
self.units = units
14+
self.abilities = abilities
15+
1016
from collections import namedtuple
11-
Build = namedtuple('Build',['id','units','abilities'])
1217
UnitRow = namedtuple('UnitRow',['id','type','title'])
1318
AbilRow = namedtuple('AbilRow',['id','type','title'])
1419

@@ -121,58 +126,209 @@
121126
'ultraliskcavern':[200,200,0],
122127
}}
123128

129+
# There are known to be duplicate ability names now. Hrm
130+
train_commands = {
131+
'RavenBuildPointDefenseDrone': 'PointDefenseDrone',
132+
'CalldownMULE': 'MULE',
133+
'BuildCommandCenter': 'CommandCenter',
134+
'BuildSupplyDepot': 'SupplyDepot',
135+
'BuildRefinery': 'Refinery',
136+
'BuildBarracks': 'Barracks',
137+
'BuildEngineeringBay': 'EngineeringBay',
138+
'BuildMissileTurret': 'MissileTurret',
139+
'BuildBunker': 'Bunker',
140+
'BuildSensorTower': 'SensorTower',
141+
'BuildGhostAcademy': 'GhostAcademy',
142+
'BuildFactory': 'Factory',
143+
'BuildStarport': 'Starport',
144+
'BuildArmory': 'Armory',
145+
'BuildFusionCore': 'FusionCore',
146+
'RavenBuildAutoTurret': 'AutoTurret',
147+
'TrainSCV': 'SCV',
148+
'TrainMarine': 'Marine',
149+
'TrainReaper': 'Reaper',
150+
'TrainGhost': 'Ghost',
151+
'TrainMarauder': 'Marauder',
152+
'TrainSiegeTank': 'SiegeTank',
153+
'TrainThor': 'Thor',
154+
'TrainHellion': 'Hellion',
155+
'TrainMedivac': 'Medivac',
156+
'TrainBanshee': 'Banshee',
157+
'TrainRaven': 'Raven',
158+
'TrainBattlecruiser': 'Battlecruiser',
159+
'TrainViking': 'VikingFighter',
160+
'MorphToPlanetaryFortress': 'PlanetaryFortress',
161+
'MorphToOrbitalCommand': 'OrbitalCommand',
162+
163+
'TrainMothership': 'Mothership',
164+
'BuildNexus': 'Nexus',
165+
'BuildPylon': 'Pylon',
166+
'BuildAssimilator': 'Assimilator',
167+
'BuildGateway': 'Gateway',
168+
'BuildForge': 'Forge',
169+
'BuildFleetBeacon': 'FleetBeacon',
170+
'BuildTwilightCouncil': 'TwilightCouncil',
171+
'BuildPhotonCannon': 'PhotonCannon',
172+
'BuildStargate': 'Stargate',
173+
'BuildTemplarArchive': 'TemplarArchive',
174+
'BuildDarkShrine': 'DarkShrine',
175+
'BuildRoboticsBay': 'RoboticsBay',
176+
'BuildRoboticsFacility': 'RoboticsFacility',
177+
'BuildCyberneticsCore': 'CyberneticsCore',
178+
'TrainZealot': 'Zealot',
179+
'TrainStalker': 'Stalker',
180+
'TrainHighTemplar': 'HighTemplar',
181+
'TrainDarkTemplar': 'DarkTemplar',
182+
'TrainSentry': 'Sentry',
183+
'TrainPhoenix': 'Phoenix',
184+
'TrainCarrier': 'Carrier',
185+
'TrainVoidRay': 'VoidRay',
186+
'TrainWarpPrism': 'WarpPrism',
187+
'TrainObserver': 'Observer',
188+
'TrainColossus': 'Colossus',
189+
'TrainImmortal': 'Immortal',
190+
'TrainProbe': 'Probe',
191+
'ArmInterceptor': 'Interceptor',
192+
'WarpInZealot': 'Zealot',
193+
'WarpInStalker': 'Stalker',
194+
'WarpInHighTemplar': 'HighTemplar',
195+
'WarpInDarkTemplar': 'DarkTemplar',
196+
'WarpInSentry': 'Sentry',
197+
'MorphToWarpGate': 'WarpGate',
198+
'MergeArchon': 'Archon',
199+
'SentryHallucinationArchon': 'HallucinatedArchon',
200+
'SentryHallucinationColossus': 'HallucinatedColossus',
201+
'SentryHallucinationHighTemplar': 'HallucinatedHighTemplar',
202+
'SentryHallucinationImmortal': 'HallucinatedImmortal',
203+
'SentryHallucinationPhoenix': 'HallucinatedPhoenix',
204+
'SentryHallucinationProbe': 'HallucinatedProbe',
205+
'SentryHallucinationStalker': 'HallucinatedStalker',
206+
'SentryHallucinationVoidRay': 'HallucinatedVoidRay',
207+
'SentryHallucinationWarpPrism': 'HallucinatedWarpPrism',
208+
'SentryHallucinationZealot': 'HallucinatedZealot',
209+
#'MorphToGateway': 'Gateway' #not safe for now, double counting
210+
211+
'MorphToBaneling': 'Baneling',
212+
'BuildHatchery': 'Hatchery',
213+
'BuildExtractor': 'Extractor',
214+
'BuildSpawningPool': 'SpawningPool',
215+
'BuildEvolutionChamber': 'EvolutionChamber',
216+
'BuildHydraliskDen': 'HydraliskDen',
217+
'BuildSpire': 'Spire',
218+
'BuildUltraliskCavern': 'UltraliskCavern',
219+
'BuildInfestationPit': 'InfestationPit',
220+
'BuildNydusNetwork': 'NydusNetwork',
221+
'BuildBanelingNest': 'BanelingNest',
222+
'BuildRoachWarren': 'RoachWarren',
223+
'BuildSpineCrawler': 'SpineCrawler',
224+
'BuildSporeCrawler': 'SporeCrawler',
225+
'MorphToLair': 'Lair',
226+
'MorphToHive': 'Hive',
227+
'MorphToGreaterSpire': 'GreaterSpire',
228+
'TrainDrone': 'Drone',
229+
'TrainZergling': 'Zergling',
230+
'TrainOverlord': 'Overlord',
231+
'TrainHydralisk': 'Hydralisk',
232+
'TrainMutalisk': 'Mutalisk',
233+
'TrainUltralisk': 'Ultralisk',
234+
'TrainRoach': 'Roach',
235+
'TrainInfestor': 'Infestor',
236+
'TrainCorruptor': 'Corruptor',
237+
'MorphToBroodLord': 'BroodLord',
238+
'MorphToOverseer': 'Overseer',
239+
'TrainQueen': 'Queen',
240+
'QueenBuildCreepTumor': 'CreepTumor',
241+
'BuildNydusCanal': 'NydusCanal',
242+
'OverseerSpawnChangeling': 'Changeling',
243+
'InfestorSpawnInfestedTerran': 'InfestorTerran',
244+
}
245+
124246
class Unit(object):
125-
def __init__(self, unit_id):
126-
self.id = unit_id
247+
def __init__(self, unit_id):
248+
self.id = unit_id
127249

128250
class Ability(object):
129-
pass
251+
pass
130252

131253
def create_build(build):
132-
units_file = path.join(BASE_PATH, "{}_{}.csv".format(build,"units"))
133-
abils_file = path.join(BASE_PATH, "{}_{}.csv".format(build,"abilities"))
134-
with open(units_file, 'r') as data_file:
135-
units = dict()
136-
for row in [UnitRow(*line.strip().split('|')[1:]) for line in data_file]:
137-
unit_id = int(row.id, 10) << 8 | 1
138-
race, minerals, vespene, supply = "Neutral", 0, 0, 0
139-
for race in ('Protoss','Terran','Zerg'):
140-
if row.type.lower() in unit_lookup[race]:
141-
minerals, vespene, supply = unit_lookup[race][row.type.lower()]
142-
break
143-
144-
units[unit_id] = type(row.title,(Unit,), dict(
145-
type=unit_id,
146-
name=row.title,
147-
title=row.title,
148-
race=race,
149-
minerals=minerals,
150-
vespene=vespene,
151-
supply=supply,
152-
))
153-
154-
155-
with open(abils_file, 'r') as data_file:
156-
abilities = {0:type('RightClick',(Ability,), dict(type=0, name='RightClick', title='Right Click'))}
157-
for row in [line.strip().split('|') for line in data_file]:
158-
base = int(row[1],10) << 5
159-
real_abils = [(base|i,t) for i,t in enumerate(row[3:]) if t!='']
160-
for abil_id, title in real_abils:
161-
abilities[abil_id] = type(title,(Ability,), dict(
162-
type=abil_id,
163-
name=title,
164-
title=title,
165-
))
166-
167-
# Some abilities have missing entries..
168-
if len(real_abils) == 0:
169-
abilities[base] = type(row[2],(Ability,), dict(
170-
type=base,
171-
name=row[2],
172-
title=row[2],
173-
))
174-
175-
return Build(build, units, abilities)
254+
units_file = path.join(BASE_PATH, "{}_{}.csv".format(build,"units"))
255+
abils_file = path.join(BASE_PATH, "{}_{}.csv".format(build,"abilities"))
256+
with open(units_file, 'r') as data_file:
257+
units = dict()
258+
for row in [UnitRow(*line.strip().split('|')[1:]) for line in data_file]:
259+
unit_id = int(row.id, 10) << 8 | 1
260+
race, minerals, vespene, supply = "Neutral", 0, 0, 0
261+
for race in ('Protoss','Terran','Zerg'):
262+
if row.type.lower() in unit_lookup[race]:
263+
minerals, vespene, supply = unit_lookup[race][row.type.lower()]
264+
break
265+
266+
units[unit_id] = type(row.title,(Unit,), dict(
267+
type=unit_id,
268+
name=row.title,
269+
title=row.title,
270+
race=race,
271+
minerals=minerals,
272+
vespene=vespene,
273+
supply=supply,
274+
))
275+
276+
if row.title.lower() in ('probe','zealot','stalker','immortal','phoenix','hightemplar','warpprism','archon','colossus','voidray'):
277+
units[unit_id+1] = type("Hallucinated"+row.title,(Unit,), dict(
278+
type=unit_id+1,
279+
name="Hallucinated"+row.title,
280+
title="Hallucinated"+row.title,
281+
race=race,
282+
minerals=0,
283+
vespene=0,
284+
supply=0,
285+
))
286+
287+
288+
with open(abils_file, 'r') as data_file:
289+
abilities = {0:type('RightClick',(Ability,), dict(type=0, name='RightClick', title='Right Click'))}
290+
for row in [line.strip().split('|') for line in data_file]:
291+
base = int(row[1],10) << 5
292+
if base == 0: continue
293+
294+
# Temporary Hack here.
295+
if base == 0xe80:
296+
real_abils = [(0xe80,"QueueCancel0"), (0xe81,"QueueCancel1")]
297+
else:
298+
real_abils = [(base|i,t) for i,t in enumerate(row[3:]) if t.strip()!='']
299+
300+
for abil_id, title in real_abils:
301+
abilities[abil_id] = type(title,(Ability,), dict(
302+
type=abil_id,
303+
name=title,
304+
title=title,
305+
))
306+
307+
308+
# Some abilities have missing entries..
309+
if len(real_abils) == 0:
310+
abilities[base] = type(row[2],(Ability,), dict(
311+
type=base,
312+
name=row[2],
313+
title=row[2],
314+
))
315+
316+
if int(row[1],10) == 249 and build==22612:
317+
pass
318+
#print row
319+
#print abilities[0x1f20], abilities[0x1f21], abilities[0x1f22], abilities[0x1f23]
320+
321+
data = Build(build, units, abilities)
322+
for unit in units.values():
323+
setattr(data, unit.name, unit)
324+
for ability in abilities.values():
325+
if ability.name in train_commands:
326+
ability.is_build = True
327+
ability.build_unit = getattr(data,train_commands[ability.name])
328+
setattr(data, ability.name, ability)
329+
330+
return data
331+
176332

177333
build16117 = create_build(16117)
178334
build17811 = create_build(17811)

‎sc2reader/events.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -156,13 +156,13 @@ def load_context(self, replay):
156156
for player in replay.players:
157157
self.logger.error("\t"+str(player))
158158

159-
print [hex(key) for key in replay.datapack.abilities.keys()]
160-
159+
#print [hex(key) for key in sorted(replay.datapack.abilities.keys())]
161160
self.logger.error("{0}\t{1}\tMissing ability {2} from {3}".format(self.frame, self.player.name, hex(self.ability_code), replay.datapack.__class__.__name__))
162161
print "{0}\t{1}\tMissing ability {2} from {3}".format(self.frame, self.player.name, hex(self.ability_code), replay.datapack.id)
163162

164163
else:
165-
self.ability_name = replay.datapack.abilities[self.ability_code].name
164+
self.ability = replay.datapack.abilities[self.ability_code]
165+
self.ability_name = self.ability.name
166166

167167

168168
def __str__(self):
@@ -191,8 +191,13 @@ def __init__(self, frame, pid, event_type, ability, target, player, team, locati
191191
def load_context(self, replay):
192192
super(TargetAbilityEvent, self).load_context(replay)
193193

194+
""" Disabled since this seems to have gone out of bounds
195+
sc2reader/ggtracker/204927.SC2Replay
194196
if self.target_owner_id:
197+
print replay.people
198+
print self.target_owner_id
195199
self.target_owner = replay.player[self.target_owner_id]
200+
"""
196201

197202
""" Disabled since team seems to always be the same player
198203
if self.target_team_id:
@@ -278,7 +283,8 @@ def load_context(self, replay):
278283
if obj_type not in data.units:
279284
msg = "Unit Type {0} not found in {1}"
280285
self.logger.error(msg.format(hex(obj_type), data.__class__.__name__))
281-
objects.append(Unit(0x0))
286+
print msg.format(hex(obj_type), data.__class__.__name__)
287+
objects.append(Unit(obj_id))
282288

283289
else:
284290
if (obj_id, obj_type) not in replay.objects:

‎sc2reader/plugins/replay.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,11 @@ def APMTracker(replay):
8989
for event in filter(efilter, player.events):
9090
player.aps[event.second] += 1
9191
player.apm[event.second/60] += 1
92-
player.avg_apm = sum(player.apm.values())/float(len(player.apm.keys()))
92+
if len(player.apm.keys()) > 0:
93+
player.avg_apm = sum(player.apm.values())/float(len(player.apm.keys()))
94+
else:
95+
player.avg_apm = 0
96+
9397

9498
@plugin
9599
def SelectionTracker(replay):

‎sc2reader/readers.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -239,13 +239,14 @@ def __call__(self, data, replay):
239239

240240
try:
241241
while event_start != data_length:
242+
event = None
242243
fstamp += read_timestamp()
243244
pid = read_bits(5)
244245
event_type = read_bits(7)
245246

246247
# Check for a lookup
247248
if event_type in EVENT_DISPATCH:
248-
event = EVENT_DISPATCH[event_type](data, fstamp, pid, event_type)
249+
event = EVENT_DISPATCH[event_type](data, fstamp, pid, event_type)
249250

250251
# Otherwise maybe it is an unknown chunk
251252
elif event_type == 0x26:
@@ -257,6 +258,8 @@ def __call__(self, data, replay):
257258
elif event_type == 0x38:
258259
arr1 = [read_bits(32) for i in range(read_bits(8))]
259260
arr2 = [read_bits(32) for i in range(read_bits(8))]
261+
elif event_type == 0x3c:
262+
read_bytes(2)
260263
elif event_type == 0x47:
261264
read_bytes(4)
262265
elif event_type == 0x4C:
@@ -273,9 +276,7 @@ def __call__(self, data, replay):
273276
if event:
274277
if debug:
275278
event.bytes = data.read_range(event_start, data.tell())
276-
277279
append(event)
278-
event = None
279280

280281
event_start = data.tell()
281282

‎sc2reader/resources.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -426,14 +426,15 @@ def load_players(self):
426426
if 'replay.message.events' in self.raw_data:
427427
# Figure out recorder
428428
self.packets = self.raw_data['replay.message.events'].packets
429-
packet_senders = map(lambda p: p.pid, self.packets)
429+
packet_senders = set(map(lambda p: p.pid, self.packets))
430430
human_pids = map(lambda p: p.pid, self.humans)
431431
recorders = list(set(human_pids) - set(packet_senders))
432432
if len(recorders) == 1:
433433
self.recorder = self.person[recorders[0]]
434434
self.recorder.recorder = True
435435
else:
436-
raise ValueError("Get Recorder algorithm is broken!")
436+
self.recorder = None
437+
self.logger.error("{} possible recorders remain: {}".format(len(recorders), recorders))
437438

438439
player_names = sorted(map(lambda p: p.name, self.people))
439440
hash_input = self.gateway+":"+','.join(player_names)

‎sc2reader/utils.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -464,8 +464,8 @@ def extract_data_file(data_file, archive):
464464
if data_file == 'replay.message.events':
465465
try:
466466
file_data = archive.read_file(data_file, force_decompress=True)
467-
except IndexError as e:
468-
if str(e) == "string index out of range":
467+
except Exception as e:
468+
if str(e) in ("string index out of range", "Unsupported compression type."):
469469
file_data = archive.read_file(data_file, force_decompress=False)
470470
else:
471471
raise

0 commit comments

Comments
 (0)