async def on_unit_created(self, unit: Unit) -> None:
"""
Add unit to dictionaries and determine what should happen to each spawned unit.
Note: This function is called automatically.
Args:
unit (Unit): the unit created
Returns:
None
"""
self.unit_dict[unit.tag] = unit.type_id
# drone protocol (prioritize gas -> minerals)
if unit.type_id in {UnitTypeId.DRONE}:
for extractor in self.gas_buildings:
if extractor.assigned_harvesters < 3:
self.do(unit.gather(extractor))
return
for base in self.townhalls:
if base.assigned_harvesters < 16:
self.do(
unit.gather(
self.mineral_field.closest_to(base.position)))
return
# queen protocol (inject > creep > unassigned)
if unit.type_id in {UnitTypeId.QUEEN}:
if len(self.inject_queens) < min(len(self.townhalls), 3):
self.inject_queens.add(unit.tag)
return
elif len(self.creep_queens) < 4:
self.creep_queens.add(unit.tag)
return
def go_mine(self, unit: Unit):
self.knowledge.roles.clear_task(unit)
if len(self.mineral_fields) > 0:
# Go to mine in this zone
mf = self.mineral_fields[0]
self.ai.do(unit.gather(mf))
elif self.ai.townhalls.exists and self.ai.mineral_field.exists:
closest_base = self.ai.townhalls.closest_to(self.center_location)
# Go to mine in some other base
mf = self.ai.mineral_field.closest_to(closest_base)
self.ai.do(unit.gather(mf))
def assign_to_work(self, worker: Unit, work: Unit):
if worker.has_buff(BuffId.ORACLESTASISTRAPTARGET):
return # Worker is in stasis and cannot move
self.roles.set_task(UnitTask.Gathering, worker)
townhalls = self.ai.townhalls.ready
if worker.is_carrying_resource and townhalls:
closest = townhalls.closest_to(worker)
self.do(worker(AbilityId.SMART, closest))
self.do(worker.gather(work, queue=True))
else:
self.do(worker.gather(work))
def assign_probe(self, bot: sc2.BotAI, unit: Unit, prio_gas: bool):
"""
Assigns the unit (PROBE) to either the closes gas or mineral patch depending
on prio_gas
"""
if prio_gas:
if not self.assign_closest_gas(
bot, unit) and not self.assign_closest_mineral(bot, unit):
closest_mineral = bot.mineral_field.closest_to(unit)
bot.do(unit.gather(closest_mineral))
else:
if not self.assign_closest_mineral(
bot, unit) and not self.assign_closest_gas(bot, unit):
closest_mineral = bot.mineral_field.closest_to(unit)
bot.do(unit.gather(closest_mineral))
def assign_closest_gas(self, bot: sc2.BotAI, unit: Unit):
gas_buildings = bot.gas_buildings.ready.sorted_by_distance_to(
unit.position)
for mining_place in gas_buildings:
if mining_place.has_vespene and mining_place.surplus_harvesters < 0:
bot.do(unit.gather(mining_place))
return True
return False
def assign_closest_mineral(self, bot: sc2.BotAI, unit: Unit) -> bool:
"""
Assigns unit to the closest mineral field that is not utilized
"""
bases = bot.townhalls.ready.sorted_by_distance_to(unit.position)
for mining_place in bases:
if mining_place.surplus_harvesters < 0:
local_minerals = {
mineral for mineral \
in bot.mineral_field if mineral.distance_to(mining_place) <= 8
}
target_mineral = max(
local_minerals,
key=lambda mineral: mineral.mineral_contents,
default=None)
bot.do(unit.gather(target_mineral))
return True
return False
async def build_gas(self, worker: Unit):
if self.best_gas is not None and self.knowledge.can_afford(
self.unit_type):
target = self.best_gas
if not self.set_worker(worker):
return False
self.builder_tag = worker.tag
cmd = worker.build_gas(target, queue=self.has_build_order(worker))
self.do(cmd)
if self.ai.race == Race.Protoss:
# Protoss only do something else after starting gas
mf = self.ai.mineral_field.closest_to(worker)
self.ai.do(worker.gather(mf, queue=True))
self.print(f"Building {self.unit_type.name} to {target.position}")
return False