diff --git a/module/campaign/campaign_ui.py b/module/campaign/campaign_ui.py index d05e7a67a..0c8872bef 100644 --- a/module/campaign/campaign_ui.py +++ b/module/campaign/campaign_ui.py @@ -66,7 +66,10 @@ class CampaignUI(UI, CampaignOcr): if name not in self.stage_entrance: logger.warning(f'Stage not found: {name}') raise CampaignNameError - return self.stage_entrance[name] + + entrance = self.stage_entrance[name] + entrance.name = name + return entrance def ensure_campaign_ui(self, name, mode='normal'): """ diff --git a/module/map/grid_info.py b/module/map/grid_info.py index 76c917296..31a0c88d0 100644 --- a/module/map/grid_info.py +++ b/module/map/grid_info.py @@ -211,3 +211,16 @@ class GridInfo: self.is_current_fleet = False self.is_submarine = False self.is_cleared = False + + def covered_grid(self): + """Relative coordinate of the covered grid. + + Returns: + list[tuple]: + """ + if self.is_current_fleet: + return [(0, -1), (0, -2)] + if self.is_fleet or self.is_siren or self.is_mystery: + return [(0, -1)] + + return [] diff --git a/module/map/grids.py b/module/map/grids.py index a308e8402..ab63fbe20 100644 --- a/module/map/grids.py +++ b/module/map/grids.py @@ -104,6 +104,7 @@ class Grids(Perspective): image = self._image_clear_ui(image) self.image = image for grid in self: + grid.reset() grid.image = image def _image_clear_ui(self, image): diff --git a/module/map/map_base.py b/module/map/map_base.py index af97aced3..bb303b49f 100644 --- a/module/map/map_base.py +++ b/module/map/map_base.py @@ -408,14 +408,7 @@ class CampaignMap: missing[attr] -= 1 for grid in self: - if not grid.is_fleet and not grid.is_mystery and not grid.is_siren: - continue - - cover = [(0, -1)] - if grid.is_current_fleet: - cover.append((0, -2)) - - for upper in cover: + for upper in grid.covered_grid(): upper = tuple(np.array(grid.location) + upper) if upper in self: upper = self[upper] @@ -451,14 +444,7 @@ class CampaignMap: # predict for grid in self: - if not grid.is_fleet and not grid.is_mystery: - continue - - cover = [(0, -1)] - if grid.is_current_fleet: - cover.append((0, -2)) - - for upper in cover: + for upper in grid.covered_grid(): upper = tuple(np.array(grid.location) + upper) if upper in self: upper = self[upper]