replay_filtering.py

from pathlib import Path
from zephyrus_sc2_parser import parse_replay

# change this to match the directory you're using
replays = Path('replays')

# change these to desired cutoff times
BUILDING_SECONDS_INTO_GAME = 0
UPGRADE_SECONDS_INTO_GAME = 0

# change this to the minimum buildings required
BUILDING_CONDITIONS = {
  'WarpGate': 4,
  # 'Nexus': 2,
}

filtered_replays = []

def check_replay(replay):
  # get races
  races = []
  winner_race = None
  for player in replay.players.values():
    races.append(player.race)
    if replay.metadata['winner'] == player.player_id:
      winner_race = player.race
   
  # not PvT so we don't care
  # can add: or winner_race != 'Terran', if only interested in Terran wins
  if not ('Protoss' in races and 'Terran' in races):
    return False
  
  # analyzing buildings/units
  building_count = {}
  for player in replay.players.values():
    if player.race == 'Protoss':
      for obj in player.objects.values():
        # seconds * 22.4 since everything is in gameloops
        if obj.name == 'WarpGate' and obj.birth_time < (BUILDING_SECONDS_INTO_GAME * 22.4):
          if obj.name not in building_count:
            building_count[obj.name] = 0
          building_count[obj.name] += 1

#   # analyzing upgrades
#   upg_condition = False
#   for upg in player.upgrades:

#     # I think it's just called 'Charge'
#     if upg.name == 'Charge' and upg.completed_at < UPGRADE_SECONDS_INTO_GAME:
#       upg_condition = True

#   if not upg_condition:
#     return False

  for building, count in building_count.items():
    # may want to change to exact match
    if count < BUILDING_CONDITIONS[building]:
      return False
   return True
          
def recurse(dir_path):
    """
    Recursively searches directories to parse replay files
    """
    for obj_path in dir_path.iterdir():
        if obj_path.is_file():
            replay = parse_replay(obj_path)
            result = check_replay(replay)
            if result:
              filtered_replays.append(obj_path.resolve())
        elif obj_path.is_dir():
            recurse(obj_path)


recurse(replays)

# print file paths
for r in filtered_replays:
  print(r)