Make forcefield, kd8charge and parasitic bomb dummy effects

BurnySc2 · Jun 8, 2019 · 034b71a · 034b71a
1 parent 9caa515
commit 034b71a
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git a/sc2/bot_ai.py b/sc2/bot_ai.py
@@ -6,13 +6,13 @@
 from typing import Any, Dict, List, Optional, Set, Tuple, Union  # mypy type checking
 
 from .cache import property_cache_forever, property_cache_once_per_frame
-from .constants import abilityid_to_unittypeid, geyser_ids, mineral_ids
+from .constants import FakeEffektID, abilityid_to_unittypeid, geyser_ids, mineral_ids
 from .data import ActionResult, Alert, Race, Result, Target, race_gas, race_townhalls, race_worker
 from .distances import DistanceCalculation
 from .game_data import AbilityData, GameData
 
 # imports for mypy and pycharm autocomplete
-from .game_state import Blip, GameState
+from .game_state import Blip, EffectData, GameState
 from .ids.ability_id import AbilityId
 from .ids.unit_typeid import UnitTypeId
 from .ids.upgrade_id import UpgradeId
@@ -836,12 +836,16 @@ def _prepare_units(self):
             if unit.is_blip:
                 self.blips.add(Blip(unit))
             else:
+                unit_type = unit.unit_type
+                # convert these units to effects: reaper grenade, parasitic bomb dummy, forcefield
+                if unit_type in FakeEffektID:
+                    self.state.effects.add(EffectData(unit, fake=True))
+                    continue
                 unit_obj = Unit(unit, self)
                 self.all_units.append(unit_obj)
                 alliance = unit.alliance
                 # Alliance.Neutral.value = 3
                 if alliance == 3:
-                    unit_type = unit.unit_type
                     # XELNAGATOWER = 149
                     if unit_type == 149:
                         self.watchtowers.append(unit_obj)

diff --git a/sc2/constants.py b/sc2/constants.py
@@ -217,3 +217,13 @@
 }
 UNIT_PHOTONCANNON = UnitTypeId.PHOTONCANNON
 UNIT_COLOSSUS = UnitTypeId.COLOSSUS
+FakeEffektRadii = {
+    UnitTypeId.KD8CHARGE.value: 2,
+    UnitTypeId.PARASITICBOMBDUMMY.value: 3,
+    UnitTypeId.FORCEFIELD.value: 1.5,
+}
+FakeEffektID = {
+    UnitTypeId.KD8CHARGE.value: "KD8CHARGE",
+    UnitTypeId.PARASITICBOMBDUMMY.value: "PARASITICBOMB",
+    UnitTypeId.FORCEFIELD.value: "FORCEFIELD",
+}
diff --git a/sc2/game_state.py b/sc2/game_state.py
@@ -1,7 +1,9 @@
 from typing import Any, Dict, List, Optional, Set, Tuple, Union  # mypy type checking
 
+from .constants import FakeEffektID, FakeEffektRadii
 from .data import Alliance, DisplayType
 from .ids.effect_id import EffectId
+from .ids.unit_typeid import UnitTypeId
 from .ids.upgrade_id import UpgradeId
 from .pixel_map import PixelMap
 from .position import Point2, Point3
@@ -73,16 +75,23 @@ def __getattr__(self, attr):
 
 
 class EffectData:
-    def __init__(self, proto):
+    def __init__(self, proto, fake=False):
         self._proto = proto
+        self.fake = fake
 
     @property
     def id(self) -> EffectId:
-        return EffectId(self._proto.effect_id)
+        if self.fake:
+            return FakeEffektID[self._proto.unit_type]
+        else:
+            return EffectId(self._proto.effect_id)
 
     @property
     def positions(self) -> Set[Point2]:
-        return {Point2.from_proto(p) for p in self._proto.pos}
+        if self.fake:
+            return {Point2.from_proto(self._proto.pos)}
+        else:
+            return {Point2.from_proto(p) for p in self._proto.pos}
 
     @property
     def alliance(self) -> Alliance:
@@ -94,7 +103,10 @@ def owner(self) -> int:
 
     @property
     def radius(self) -> float:
-        return self._proto.radius
+        if self.fake:
+            return FakeEffektRadii[self._proto.unit_type]
+        else:
+            return self._proto.radius
 
     def __repr__(self) -> str:
         return f"{self.id} with radius {self.radius} at {self.positions}"