def __init__(self, env: StarCraft2Env, mode: str):
os.environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "hide"
self.env = env
self.mode = mode
self.obs = None
self._window_scale = 0.75
self.game_info = game_info = self.env._controller.game_info()
self.static_data = self.env._controller.data()
self._obs_queue: queue.Queue = queue.Queue()
self._game_times: collections.deque = collections.deque(maxlen=100) # Avg FPS over 100 frames. # pytype: disable=wrong-keyword-args
self._render_times: collections.deque = collections.deque(maxlen=100) # pytype: disable=wrong-keyword-args
self._last_time = time.time()
self._last_game_loop = 0
self._name_lengths: dict = {}
self._map_size = point.Point.build(game_info.start_raw.map_size)
self._playable = point.Rect(
point.Point.build(game_info.start_raw.playable_area.p0),
point.Point.build(game_info.start_raw.playable_area.p1),
)
window_size_px = point.Point(self.env.window_size[0],
self.env.window_size[1])
window_size_px = self._map_size.scale_max_size(
window_size_px * self._window_scale).ceil()
self._scale = window_size_px.y // 32
self.display = pygame.Surface(window_size_px)
if mode == "human":
self.display = pygame.display.set_mode(
window_size_px, 0, 32) # type: ignore # noqa: E501
pygame.display.init()
pygame.display.set_caption("Starcraft Viewer")
pygame.font.init()
self._world_to_world_tl = transform.Linear(
point.Point(1, -1), point.Point(0, self._map_size.y))
self._world_tl_to_screen = transform.Linear(scale=window_size_px / 32)
self.screen_transform = transform.Chain(self._world_to_world_tl,
self._world_tl_to_screen)
surf_loc = point.Rect(point.origin, window_size_px)
sub_surf = self.display.subsurface(
pygame.Rect(surf_loc.tl, surf_loc.size))
self._surf = _Surface(sub_surf, None, surf_loc, self.screen_transform,
None, self.draw_screen)
self._font_small = pygame.font.Font(None, int(self._scale * 0.5))
self._font_large = pygame.font.Font(None, self._scale)
self.upgrade_colors = [
colors.black, # unused...
colors.white * 0.6,
colors.white * 0.8,
colors.white,
]
def testInitBad(self):
with self.assertRaises(TypeError):
point.Rect(4, 3, 2, 1) # require t <= b, l <= r
with self.assertRaises(TypeError):
point.Rect(1)
with self.assertRaises(TypeError):
point.Rect(1, 2, 3)
with self.assertRaises(TypeError):
point.Rect()
def draw_actions(self):
"""Draw the actions so that they can be inspected for accuracy."""
for act in self._obs.actions:
if (act.HasField("action_raw") and
act.action_raw.HasField("unit_command") and
act.action_raw.unit_command.HasField("target_world_space_pos")):
pos = point.Point.build(
act.action_raw.unit_command.target_world_space_pos)
self.all_surfs(_Surface.draw_circle, colors.yellow, pos, 0.1)
if act.HasField("action_feature_layer"):
act_fl = act.action_feature_layer
if act_fl.HasField("unit_command"):
if act_fl.unit_command.HasField("target_screen_coord"):
pos = self._world_to_feature_screen_px.back_pt(
point.Point.build(act_fl.unit_command.target_screen_coord))
self.all_surfs(_Surface.draw_circle, colors.cyan, pos, 0.1)
if act_fl.unit_command.HasField("target_minimap_coord"):
pos = self._world_to_feature_minimap_px.back_pt(
point.Point.build(act_fl.unit_command.target_minimap_coord))
self.all_surfs(_Surface.draw_circle, colors.cyan, pos, 0.1)
if (act_fl.HasField("unit_selection_point") and
act_fl.unit_selection_point.HasField("selection_screen_coord")):
pos = self._world_to_feature_screen_px.back_pt(point.Point.build(
act_fl.unit_selection_point.selection_screen_coord))
self.all_surfs(_Surface.draw_circle, colors.cyan, pos, 0.1)
if act_fl.HasField("unit_selection_rect"):
for r in act_fl.unit_selection_rect.selection_screen_coord:
rect = point.Rect(
self._world_to_feature_screen_px.back_pt(
point.Point.build(r.p0)),
self._world_to_feature_screen_px.back_pt(
point.Point.build(r.p1)))
self.all_surfs(_Surface.draw_rect, colors.cyan, rect, 1)
if act.HasField("action_render"):
act_rgb = act.action_render
if act_rgb.HasField("unit_command"):
if act_rgb.unit_command.HasField("target_screen_coord"):
pos = self._world_to_rgb_screen_px.back_pt(
point.Point.build(act_rgb.unit_command.target_screen_coord))
self.all_surfs(_Surface.draw_circle, colors.red, pos, 0.1)
if act_rgb.unit_command.HasField("target_minimap_coord"):
pos = self._world_to_rgb_minimap_px.back_pt(
point.Point.build(act_rgb.unit_command.target_minimap_coord))
self.all_surfs(_Surface.draw_circle, colors.red, pos, 0.1)
if (act_rgb.HasField("unit_selection_point") and
act_rgb.unit_selection_point.HasField("selection_screen_coord")):
pos = self._world_to_rgb_screen_px.back_pt(point.Point.build(
act_rgb.unit_selection_point.selection_screen_coord))
self.all_surfs(_Surface.draw_circle, colors.red, pos, 0.1)
if act_rgb.HasField("unit_selection_rect"):
for r in act_rgb.unit_selection_rect.selection_screen_coord:
rect = point.Rect(
self._world_to_rgb_screen_px.back_pt(
point.Point.build(r.p0)),
self._world_to_rgb_screen_px.back_pt(
point.Point.build(r.p1)))
self.all_surfs(_Surface.draw_rect, colors.red, rect, 1)
def draw_selection(self, surf):
"""Draw the selection rectange."""
if self.select_start:
mouse_pos = self.get_mouse_pos()
if mouse_pos and mouse_pos.type == SurfType.SCREEN:
surf.draw_rect(
colors.green, point.Rect(self.select_start, mouse_pos.pos), 1)
def world_to_screen_pos(game_info, pos, obs):
"""
:param game_info: env.game_info
:param pos: target_world_space_pos
:param obs: obs.raw_observation.observation.raw_data.player.camera
:return: screen_pos
"""
# init parameter and define
map_size = point.Point.build(game_info.start_raw.map_size)
fl_opts = game_info.options.feature_layer
feature_layer_screen_size = point.Point.build(fl_opts.resolution)
camera_width_world_units = fl_opts.width
world_to_screen = transform.Linear(point.Point(1, -1),
point.Point(0, map_size.y))
screen_to_fl_screen = transform.Linear(feature_layer_screen_size /
camera_width_world_units)
world_to_fl_screen = transform.Chain(world_to_screen, screen_to_fl_screen,
transform.Floor())
# Update the camera transform based on the new camera center.
camera_center = obs.raw_observation.observation.raw_data.player.camera
camera_radius = (feature_layer_screen_size / feature_layer_screen_size.x *
camera_width_world_units / 2)
camera_center = point.Point.build(camera_center)
center = camera_center.bound(camera_radius, map_size - camera_radius)
camera = point.Rect((center - camera_radius).bound(map_size),
(center + camera_radius).bound(map_size))
world_to_screen.offset = (-camera.bl * world_to_screen.scale)
trans_pos = world_to_fl_screen.fwd_pt(point.Point.build(pos))
return np.clip(np.array(trans_pos), 0, 63).tolist()
def select_action(self, pos1, pos2, ctrl, shift):
"""Return a `sc_pb.Action` with the selection filled."""
assert pos1.surf.surf_type == pos2.surf.surf_type
assert pos1.surf.world_to_obs == pos2.surf.world_to_obs
action = sc_pb.Action()
action_spatial = pos1.action_spatial(action)
if pos1.world_pos == pos2.world_pos: # select a point
select = action_spatial.unit_selection_point
pos1.obs_pos.assign_to(select.selection_screen_coord)
mod = sc_spatial.ActionSpatialUnitSelectionPoint
if ctrl:
select.type = mod.AddAllType if shift else mod.AllType
else:
select.type = mod.Toggle if shift else mod.Select
else:
select = action_spatial.unit_selection_rect
rect = select.selection_screen_coord.add()
pos1.obs_pos.assign_to(rect.p0)
pos2.obs_pos.assign_to(rect.p1)
select.selection_add = shift
# Clear the queued action if something will be selected. An alternative
# implementation may check whether the selection changed next frame.
units = self._units_in_area(point.Rect(pos1.world_pos, pos2.world_pos))
if units:
self.clear_queued_action()
return action
def draw_selection(self, surf):
"""Draw the selection rectange."""
if self._select_start:
mouse_pos = self.get_mouse_pos()
if (mouse_pos and mouse_pos.surf.surf_type & SurfType.SCREEN and
mouse_pos.surf.surf_type == self._select_start.surf.surf_type):
rect = point.Rect(self._select_start.world_pos, mouse_pos.world_pos)
surf.draw_rect(colors.green, rect, 1)
def testInit(self):
r = point.Rect(1, 2, 3, 4)
self.assertEqual(r.t, 1)
self.assertEqual(r.l, 2)
self.assertEqual(r.b, 3)
self.assertEqual(r.r, 4)
self.assertEqual(r.tl, point.Point(2, 1))
self.assertEqual(r.tr, point.Point(4, 1))
self.assertEqual(r.bl, point.Point(2, 3))
self.assertEqual(r.br, point.Point(4, 3))
def _update_camera(self, camera_center):
"""Update the camera transform based on the new camera center."""
camera_radius = (self._feature_layer_screen_size /
self._feature_layer_screen_size.x *
self._camera_width_world_units / 2)
center = camera_center.bound(camera_radius, self._map_size - camera_radius)
self._camera = point.Rect(
(center - camera_radius).bound(self._map_size),
(center + camera_radius).bound(self._map_size))
self._world_to_screen.offset = (-self._camera.bl *
self._world_to_screen.scale)
def testInitOnePoint(self):
r = point.Rect(point.Point(1, 2))
self.assertEqual(r.t, 0)
self.assertEqual(r.l, 0)
self.assertEqual(r.b, 2)
self.assertEqual(r.r, 1)
self.assertEqual(r.tl, point.Point(0, 0))
self.assertEqual(r.tr, point.Point(1, 0))
self.assertEqual(r.bl, point.Point(0, 2))
self.assertEqual(r.br, point.Point(1, 2))
self.assertEqual(r.size, point.Point(1, 2))
self.assertEqual(r.center, point.Point(1, 2) / 2)
self.assertEqual(r.area, 2)
def testInitTwoPointsReversed(self):
r = point.Rect(point.Point(3, 4), point.Point(1, 2))
self.assertEqual(r.t, 2)
self.assertEqual(r.l, 1)
self.assertEqual(r.b, 4)
self.assertEqual(r.r, 3)
self.assertEqual(r.tl, point.Point(1, 2))
self.assertEqual(r.tr, point.Point(3, 2))
self.assertEqual(r.bl, point.Point(1, 4))
self.assertEqual(r.br, point.Point(3, 4))
self.assertEqual(r.size, point.Point(2, 2))
self.assertEqual(r.center, point.Point(2, 3))
self.assertEqual(r.area, 4)
def _update_camera(self, camera_center):
"""Update the camera transform based on the new camera center."""
self._world_tl_to_world_camera_rel.offset = (
-self._world_to_world_tl.fwd_pt(camera_center) *
self._world_tl_to_world_camera_rel.scale)
if self._feature_screen_px:
camera_radius = (self._feature_screen_px / self._feature_screen_px.x *
self._feature_camera_width_world_units / 2)
center = camera_center.bound(camera_radius,
self._map_size - camera_radius)
self._camera = point.Rect(
(center - camera_radius).bound(self._map_size),
(center + camera_radius).bound(self._map_size))
def add_feature_layer(feature, surf_type, world_to_surf):
"""Add a feature layer surface."""
i = next(feature_counter)
grid_offset = point.Point(i % cols, i // cols) * feature_grid_size
text = feature_font.render(feature.full_name, True, colors.white)
rect = text.get_rect()
rect.center = grid_offset + point.Point(feature_grid_size.x / 2,
feature_font_size)
feature_pane.blit(text, rect)
surf_loc = (features_loc + grid_offset + feature_layer_padding +
point.Point(0, feature_font_size))
add_surface(surf_type,
point.Rect(surf_loc, surf_loc + feature_layer_size),
world_to_surf,
lambda surf: self.draw_feature_layer(surf, feature))
def screen_to_minimap_pos(game_info, screen_pos, obs):
screen_pos = Pos(screen_pos[0], screen_pos[1])
# init parameter and define
map_size = point.Point.build(game_info.start_raw.map_size)
fl_opts = game_info.options.feature_layer
feature_layer_screen_size = point.Point.build(fl_opts.resolution)
feature_layer_minimap_size = point.Point.build(fl_opts.minimap_resolution)
# screen to world
camera_width_world_units = fl_opts.width
world_to_screen = transform.Linear(point.Point(1, -1),
point.Point(0, map_size.y))
screen_to_fl_screen = transform.Linear(feature_layer_screen_size /
camera_width_world_units)
world_to_fl_screen = transform.Chain(world_to_screen, screen_to_fl_screen,
transform.Floor())
# Update the camera transform based on the new camera center.
camera_center = obs.raw_observation.observation.raw_data.player.camera
camera_radius = (feature_layer_screen_size / feature_layer_screen_size.x *
camera_width_world_units / 2)
camera_center = point.Point.build(camera_center)
center = camera_center.bound(camera_radius, map_size - camera_radius)
camera = point.Rect((center - camera_radius).bound(map_size),
(center + camera_radius).bound(map_size))
world_to_screen.offset = (-camera.bl * world_to_screen.scale)
world_pos = world_to_fl_screen.back_pt(point.Point.build(screen_pos))
# world to minimap
max_map_dim = map_size.max_dim()
world_to_minimap = transform.Linear(point.Point(1, -1),
point.Point(0, map_size.y))
minimap_to_fl_minimap = transform.Linear(feature_layer_minimap_size /
max_map_dim)
world_to_fl_minimap = transform.Chain(world_to_minimap,
minimap_to_fl_minimap,
transform.Floor())
minimap_pos = world_to_fl_minimap.fwd_pt(point.Point.build(world_pos))
return np.clip(np.array(minimap_pos), 0, 63).tolist()
def rect(a):
return point.Rect(point.Point(*a[1]).floor(),
point.Point(*a[2]).floor())
def get_actions(self, run_config, controller):
"""Get actions from the UI, apply to controller, and return an ActionCmd."""
if not self._initialized:
return ActionCmd.STEP
for event in pygame.event.get():
if event.type == pygame.QUIT:
return ActionCmd.QUIT
elif event.type == pygame.KEYDOWN:
if self.help:
self.help = False
elif event.key in (pygame.K_QUESTION, pygame.K_SLASH):
self.help = True
elif event.key == pygame.K_PAUSE:
pause = True
while pause:
time.sleep(0.1)
for event in pygame.event.get():
if event.type == pygame.KEYDOWN:
if event.key in (pygame.K_PAUSE, pygame.K_ESCAPE):
pause = False
elif event.key == pygame.K_F4:
return ActionCmd.QUIT
elif event.key == pygame.K_F5:
return ActionCmd.RESTART
elif event.key == pygame.K_F4:
return ActionCmd.QUIT
elif event.key == pygame.K_F5:
return ActionCmd.RESTART
elif event.key == pygame.K_F8: # Toggle synchronous rendering.
self._render_sync = not self._render_sync
print("Rendering", self._render_sync and "Sync" or "Async")
elif event.key == pygame.K_F9: # Save a replay.
self.save_replay(run_config, controller)
elif (event.key in (pygame.K_PLUS, pygame.K_EQUALS) and
pygame.key.get_mods() & pygame.KMOD_CTRL): # zoom in
self.zoom(1.1)
elif (event.key in (pygame.K_MINUS, pygame.K_UNDERSCORE) and
pygame.key.get_mods() & pygame.KMOD_CTRL): # zoom out
self.zoom(1 / 1.1)
elif event.key in (pygame.K_PAGEUP, pygame.K_PAGEDOWN):
if pygame.key.get_mods() & pygame.KMOD_CTRL:
if event.key == pygame.K_PAGEUP:
self._step_mul += 1
elif self._step_mul > 1:
self._step_mul -= 1
print("New step mul:", self._step_mul)
else:
self._fps *= 1.25 if event.key == pygame.K_PAGEUP else 1 / 1.25
print("New max game speed: %.1f" % self._fps)
elif event.key in self.camera_actions:
controller.act(self.camera_action_raw(
self._camera.center + self.camera_actions[event.key]))
elif event.key == pygame.K_ESCAPE:
cmds = self._abilities(lambda cmd: cmd.hotkey == "escape")
if cmds:
assert len(cmds) == 1
cmd = cmds[0]
assert cmd.target == sc_data.AbilityData.Target.Value("None")
controller.act(self.unit_action(cmd))
else:
self.clear_queued_action()
else:
if not self.queued_action:
key = pygame.key.name(event.key).lower()
new_cmd = self.queued_hotkey + key
cmds = self._abilities(lambda cmd, n=new_cmd: ( # pylint: disable=g-long-lambda
cmd.hotkey != "escape" and cmd.hotkey.startswith(n)))
if cmds:
self.queued_hotkey = new_cmd
if len(cmds) == 1:
cmd = cmds[0]
if cmd.hotkey == self.queued_hotkey:
if cmd.target != sc_data.AbilityData.Target.Value("None"):
self.clear_queued_action()
self.queued_action = cmd
else:
controller.act(self.unit_action(cmd))
elif event.type == pygame.MOUSEBUTTONDOWN:
mouse_pos = self.get_mouse_pos(event.pos)
if event.button == MouseButtons.LEFT and mouse_pos:
if self.queued_action:
controller.act(self.unit_action(self.queued_action, mouse_pos))
elif mouse_pos.type == SurfType.MINIMAP:
controller.act(self.camera_action(mouse_pos.pos))
else:
self.select_start = mouse_pos.pos
elif event.button == MouseButtons.RIGHT:
if self.queued_action:
self.clear_queued_action()
else:
cmds = self._abilities(lambda cmd: cmd.button_name == "Smart")
if cmds:
controller.act(self.unit_action(cmds[0], mouse_pos))
elif event.type == pygame.MOUSEBUTTONUP:
mouse_pos = self.get_mouse_pos(event.pos)
if event.button == MouseButtons.LEFT and self.select_start:
if mouse_pos and mouse_pos.type == SurfType.SCREEN:
controller.act(self.select_action(point.Rect(self.select_start,
mouse_pos.pos)))
self.select_start = None
return ActionCmd.STEP
def init_window(self):
"""Initialize the pygame window and lay out the surfaces."""
pygame.init()
# Want a roughly square grid of feature layers, each being roughly square.
num_feature_layers = (len(features.SCREEN_FEATURES) +
len(features.MINIMAP_FEATURES))
cols = math.ceil(math.sqrt(num_feature_layers))
rows = math.ceil(num_feature_layers / cols)
features_layout = point.Point(cols, rows * 1.05) # make room for titles
# Scale such that features_layout and screen_aspect ratio have the same
# height so that we can figure out the max window size and ratio of widths.
screen_aspect_ratio = (self._feature_layer_screen_size *
(rows / self._feature_layer_screen_size.y))
total = features_layout + point.Point(screen_aspect_ratio.x, 0)
window_size_px = total.scale_max_size(_get_max_window_size()).ceil()
# Create the actual window surface. This should only be blitted to from one
# of the sub-surfaces defined below.
self._window = pygame.display.set_mode(window_size_px, 0, 32)
pygame.display.set_caption("Starcraft Viewer")
# The sub-surfaces that the various draw functions will draw to.
self.surfaces = []
def add_surface(surf_type, surf_loc, world_to_surf, draw_fn):
"""Add a surface. Drawn in order and intersect in reverse order."""
sub_surf = self._window.subsurface(
pygame.Rect(surf_loc.tl, surf_loc.size))
self.surfaces.append(_Surface(
sub_surf, surf_type, surf_loc, world_to_surf, draw_fn))
self.scale = window_size_px.y // 30
self.font_small = pygame.font.Font(None, int(self.scale * 0.5))
self.font_large = pygame.font.Font(None, self.scale)
# Just flip so the base minimap is TL origin
self._world_to_minimap = transform.Linear(point.Point(1, -1),
point.Point(0, self._map_size.y))
max_map_dim = self._map_size.max_dim()
self._minimap_to_fl_minimap = transform.Linear(
self._feature_layer_minimap_size / max_map_dim)
self._world_to_fl_minimap = transform.Chain(
self._world_to_minimap,
self._minimap_to_fl_minimap,
transform.Floor())
# Flip and zoom to the camera area. Update the offset as the camera moves.
self._world_to_screen = transform.Linear(point.Point(1, -1),
point.Point(0, self._map_size.y))
self._screen_to_fl_screen = transform.Linear(
self._feature_layer_screen_size / self._camera_width_world_units)
self._world_to_fl_screen = transform.Chain(
self._world_to_screen,
self._screen_to_fl_screen,
transform.Floor())
# Renderable space for the screen.
self.screen_size_px = self._feature_layer_screen_size.scale_max_size(
window_size_px)
screen_to_visual_screen = transform.Linear(
self.screen_size_px.x / self._camera_width_world_units)
add_surface(SurfType.SCREEN,
point.Rect(point.origin, self.screen_size_px),
transform.Chain(
self._world_to_screen,
screen_to_visual_screen),
self.draw_screen)
# Renderable space for the minimap.
self.minimap_size_px = self._map_size.scale_max_size(
self.screen_size_px / 4)
minimap_to_visual_minimap = transform.Linear(
self.minimap_size_px.max_dim() / max_map_dim)
minimap_offset = point.Point(0, (self.screen_size_px.y -
self.minimap_size_px.y))
add_surface(SurfType.MINIMAP,
point.Rect(minimap_offset,
minimap_offset + self.minimap_size_px),
transform.Chain(
self._world_to_minimap,
minimap_to_visual_minimap),
self.draw_mini_map)
# Add the feature layers
features_loc = point.Point(self.screen_size_px.x, 0)
feature_pane = self._window.subsurface(
pygame.Rect(features_loc, window_size_px - features_loc))
feature_pane.fill(colors.white / 2)
feature_pane_size = point.Point(*feature_pane.get_size())
feature_grid_size = feature_pane_size / point.Point(cols, rows)
feature_layer_area = self._feature_layer_screen_size.scale_max_size(
feature_grid_size)
feature_layer_size = feature_layer_area * 0.9
feature_layer_padding = (feature_layer_area - feature_layer_size) / 2
feature_font_size = int(feature_grid_size.y * 0.09)
feature_font = pygame.font.Font(None, feature_font_size)
feature_counter = itertools.count()
def add_feature_layer(feature, surf_type, world_to_surf):
"""Add a feature layer surface."""
i = next(feature_counter)
grid_offset = point.Point(i % cols, i // cols) * feature_grid_size
text = feature_font.render(feature.full_name, True, colors.white)
rect = text.get_rect()
rect.center = grid_offset + point.Point(feature_grid_size.x / 2,
feature_font_size)
feature_pane.blit(text, rect)
surf_loc = (features_loc + grid_offset + feature_layer_padding +
point.Point(0, feature_font_size))
add_surface(surf_type,
point.Rect(surf_loc, surf_loc + feature_layer_size),
world_to_surf,
lambda surf: self.draw_feature_layer(surf, feature))
# Add the minimap feature layers
fl_minimap_to_fl_surf = transform.Linear(
feature_layer_size / self._feature_layer_minimap_size)
world_to_fl_minimap_surf = transform.Chain(
self._world_to_minimap,
self._minimap_to_fl_minimap,
transform.Center(),
fl_minimap_to_fl_surf)
for feature in features.MINIMAP_FEATURES:
add_feature_layer(feature, SurfType.MINIMAP, world_to_fl_minimap_surf)
# Add the screen feature layers
fl_screen_to_fl_surf = transform.Linear(
feature_layer_size / self._feature_layer_screen_size)
world_to_fl_screen_surf = transform.Chain(
self._world_to_screen,
self._screen_to_fl_screen,
transform.Center(),
fl_screen_to_fl_surf)
for feature in features.SCREEN_FEATURES:
add_feature_layer(feature, SurfType.SCREEN, world_to_fl_screen_surf)
# Add the help screen
add_surface(SurfType.CHROME,
point.Rect(window_size_px / 4, window_size_px * 3 / 4),
None,
self.draw_help)
# Arbitrarily set the initial camera to the center of the map.
self._update_camera(self._map_size / 2)
def init_window(self):
"""Initialize the pygame window and lay out the surfaces."""
if os.name == "nt":
# Enable DPI awareness on Windows to give the correct window size.
ctypes.windll.user32.SetProcessDPIAware() # pytype: disable=module-attr
pygame.init()
if self._render_rgb and self._rgb_screen_px:
main_screen_px = self._rgb_screen_px
else:
main_screen_px = self._feature_screen_px
window_size_ratio = main_screen_px
if self._feature_screen_px and self._render_feature_grid:
# Want a roughly square grid of feature layers, each being roughly square.
num_feature_layers = (len(features.SCREEN_FEATURES) +
len(features.MINIMAP_FEATURES))
feature_cols = math.ceil(math.sqrt(num_feature_layers))
feature_rows = math.ceil(num_feature_layers / feature_cols)
features_layout = point.Point(feature_cols,
feature_rows * 1.05) # make room for titles
# Scale features_layout to main_screen_px height so we know its width.
features_aspect_ratio = (features_layout * main_screen_px.y /
features_layout.y)
window_size_ratio += point.Point(features_aspect_ratio.x, 0)
window_size_px = window_size_ratio.scale_max_size(
_get_max_window_size()).ceil()
# Create the actual window surface. This should only be blitted to from one
# of the sub-surfaces defined below.
self._window = pygame.display.set_mode(window_size_px, 0, 32)
pygame.display.set_caption("Starcraft Viewer")
# The sub-surfaces that the various draw functions will draw to.
self._surfaces = []
def add_surface(surf_type, surf_loc, world_to_surf, world_to_obs, draw_fn):
"""Add a surface. Drawn in order and intersect in reverse order."""
sub_surf = self._window.subsurface(
pygame.Rect(surf_loc.tl, surf_loc.size))
self._surfaces.append(_Surface(
sub_surf, surf_type, surf_loc, world_to_surf, world_to_obs, draw_fn))
self._scale = window_size_px.y // 30
self._font_small = pygame.font.Font(None, int(self._scale * 0.5))
self._font_large = pygame.font.Font(None, self._scale)
def check_eq(a, b):
"""Used to run unit tests on the transforms."""
assert (a - b).len() < 0.0001, "%s != %s" % (a, b)
# World has origin at bl, world_tl has origin at tl.
self._world_to_world_tl = transform.Linear(
point.Point(1, -1), point.Point(0, self._map_size.y))
check_eq(self._world_to_world_tl.fwd_pt(point.Point(0, 0)),
point.Point(0, self._map_size.y))
check_eq(self._world_to_world_tl.fwd_pt(point.Point(5, 10)),
point.Point(5, self._map_size.y - 10))
# Move the point to be relative to the camera. This gets updated per frame.
self._world_tl_to_world_camera_rel = transform.Linear(
offset=-self._map_size / 4)
check_eq(self._world_tl_to_world_camera_rel.fwd_pt(self._map_size / 4),
point.Point(0, 0))
check_eq(
self._world_tl_to_world_camera_rel.fwd_pt(
(self._map_size / 4) + point.Point(5, 10)),
point.Point(5, 10))
if self._feature_screen_px:
# Feature layer locations in continuous space.
feature_world_per_pixel = (self._feature_screen_px /
self._feature_camera_width_world_units)
world_camera_rel_to_feature_screen = transform.Linear(
feature_world_per_pixel, self._feature_screen_px / 2)
check_eq(world_camera_rel_to_feature_screen.fwd_pt(point.Point(0, 0)),
self._feature_screen_px / 2)
check_eq(
world_camera_rel_to_feature_screen.fwd_pt(
point.Point(-0.5, -0.5) * self._feature_camera_width_world_units),
point.Point(0, 0))
self._world_to_feature_screen = transform.Chain(
self._world_to_world_tl,
self._world_tl_to_world_camera_rel,
world_camera_rel_to_feature_screen)
self._world_to_feature_screen_px = transform.Chain(
self._world_to_feature_screen,
transform.PixelToCoord())
world_tl_to_feature_minimap = transform.Linear(
self._feature_minimap_px / self._map_size.max_dim())
check_eq(world_tl_to_feature_minimap.fwd_pt(point.Point(0, 0)),
point.Point(0, 0))
check_eq(world_tl_to_feature_minimap.fwd_pt(self._map_size),
self._map_size.scale_max_size(self._feature_minimap_px))
self._world_to_feature_minimap = transform.Chain(
self._world_to_world_tl,
world_tl_to_feature_minimap)
self._world_to_feature_minimap_px = transform.Chain(
self._world_to_feature_minimap,
transform.PixelToCoord())
if self._rgb_screen_px:
# RGB pixel locations in continuous space.
# TODO(tewalds): Use a real 3d projection instead of orthogonal.
rgb_world_per_pixel = (self._rgb_screen_px / 24)
world_camera_rel_to_rgb_screen = transform.Linear(
rgb_world_per_pixel, self._rgb_screen_px / 2)
check_eq(world_camera_rel_to_rgb_screen.fwd_pt(point.Point(0, 0)),
self._rgb_screen_px / 2)
check_eq(
world_camera_rel_to_rgb_screen.fwd_pt(
point.Point(-0.5, -0.5) * 24),
point.Point(0, 0))
self._world_to_rgb_screen = transform.Chain(
self._world_to_world_tl,
self._world_tl_to_world_camera_rel,
world_camera_rel_to_rgb_screen)
self._world_to_rgb_screen_px = transform.Chain(
self._world_to_rgb_screen,
transform.PixelToCoord())
world_tl_to_rgb_minimap = transform.Linear(
self._rgb_minimap_px / self._map_size.max_dim())
check_eq(world_tl_to_rgb_minimap.fwd_pt(point.Point(0, 0)),
point.Point(0, 0))
check_eq(world_tl_to_rgb_minimap.fwd_pt(self._map_size),
self._map_size.scale_max_size(self._rgb_minimap_px))
self._world_to_rgb_minimap = transform.Chain(
self._world_to_world_tl,
world_tl_to_rgb_minimap)
self._world_to_rgb_minimap_px = transform.Chain(
self._world_to_rgb_minimap,
transform.PixelToCoord())
# Renderable space for the screen.
screen_size_px = main_screen_px.scale_max_size(window_size_px)
minimap_size_px = self._map_size.scale_max_size(screen_size_px / 4)
minimap_offset = point.Point(0, (screen_size_px.y - minimap_size_px.y))
if self._render_rgb:
rgb_screen_to_main_screen = transform.Linear(
screen_size_px / self._rgb_screen_px)
add_surface(SurfType.RGB | SurfType.SCREEN,
point.Rect(point.origin, screen_size_px),
transform.Chain( # surf
self._world_to_rgb_screen,
rgb_screen_to_main_screen),
self._world_to_rgb_screen_px,
self.draw_screen)
rgb_minimap_to_main_minimap = transform.Linear(
minimap_size_px / self._rgb_minimap_px)
add_surface(SurfType.RGB | SurfType.MINIMAP,
point.Rect(minimap_offset,
minimap_offset + minimap_size_px),
transform.Chain( # surf
self._world_to_rgb_minimap,
rgb_minimap_to_main_minimap),
self._world_to_rgb_minimap_px,
self.draw_mini_map)
else:
feature_screen_to_main_screen = transform.Linear(
screen_size_px / self._feature_screen_px)
add_surface(SurfType.FEATURE | SurfType.SCREEN,
point.Rect(point.origin, screen_size_px),
transform.Chain( # surf
self._world_to_feature_screen,
feature_screen_to_main_screen),
self._world_to_feature_screen_px,
self.draw_screen)
feature_minimap_to_main_minimap = transform.Linear(
minimap_size_px / self._feature_minimap_px)
add_surface(SurfType.FEATURE | SurfType.MINIMAP,
point.Rect(minimap_offset,
minimap_offset + minimap_size_px),
transform.Chain( # surf
self._world_to_feature_minimap,
feature_minimap_to_main_minimap),
self._world_to_feature_minimap_px,
self.draw_mini_map)
if self._feature_screen_px and self._render_feature_grid:
# Add the feature layers
features_loc = point.Point(screen_size_px.x, 0)
feature_pane = self._window.subsurface(
pygame.Rect(features_loc, window_size_px - features_loc))
feature_pane.fill(colors.white / 2)
feature_pane_size = point.Point(*feature_pane.get_size())
feature_grid_size = feature_pane_size / point.Point(feature_cols,
feature_rows)
feature_layer_area = self._feature_screen_px.scale_max_size(
feature_grid_size)
feature_layer_padding = feature_layer_area // 20
feature_layer_size = feature_layer_area - feature_layer_padding * 2
feature_font_size = int(feature_grid_size.y * 0.09)
feature_font = pygame.font.Font(None, feature_font_size)
feature_counter = itertools.count()
def add_feature_layer(feature, surf_type, world_to_surf, world_to_obs):
"""Add a feature layer surface."""
i = next(feature_counter)
grid_offset = point.Point(i % feature_cols,
i // feature_cols) * feature_grid_size
text = feature_font.render(feature.full_name, True, colors.white)
rect = text.get_rect()
rect.center = grid_offset + point.Point(feature_grid_size.x / 2,
feature_font_size)
feature_pane.blit(text, rect)
surf_loc = (features_loc + grid_offset + feature_layer_padding +
point.Point(0, feature_font_size))
add_surface(surf_type,
point.Rect(surf_loc, surf_loc + feature_layer_size),
world_to_surf, world_to_obs,
lambda surf: self.draw_feature_layer(surf, feature))
# Add the minimap feature layers
feature_minimap_to_feature_minimap_surf = transform.Linear(
feature_layer_size / self._feature_minimap_px)
world_to_feature_minimap_surf = transform.Chain(
self._world_to_feature_minimap,
feature_minimap_to_feature_minimap_surf)
for feature in features.MINIMAP_FEATURES:
add_feature_layer(feature, SurfType.FEATURE | SurfType.MINIMAP,
world_to_feature_minimap_surf,
self._world_to_feature_minimap_px)
# Add the screen feature layers
feature_screen_to_feature_screen_surf = transform.Linear(
feature_layer_size / self._feature_screen_px)
world_to_feature_screen_surf = transform.Chain(
self._world_to_feature_screen,
feature_screen_to_feature_screen_surf)
for feature in features.SCREEN_FEATURES:
add_feature_layer(feature, SurfType.FEATURE | SurfType.SCREEN,
world_to_feature_screen_surf,
self._world_to_feature_screen_px)
# Add the help screen
help_size = point.Point(
(max(len(s) for s, _ in self.shortcuts) +
max(len(s) for _, s in self.shortcuts)) * 0.4 + 4,
len(self.shortcuts) + 3) * self._scale
help_rect = point.Rect(window_size_px / 2 - help_size / 2,
window_size_px / 2 + help_size / 2)
add_surface(SurfType.CHROME, help_rect, None, None, self.draw_help)
# Arbitrarily set the initial camera to the center of the map.
self._update_camera(self._map_size / 2)
def testIntersectsCircle(self):
r = point.Rect(point.Point(1, 1), point.Point(3, 3))
self.assertFalse(r.intersects_circle(point.Point(0, 0), 0.5))
self.assertFalse(r.intersects_circle(point.Point(0, 0), 1))
self.assertTrue(r.intersects_circle(point.Point(0, 0), 1.5))
self.assertTrue(r.intersects_circle(point.Point(0, 0), 2))
def testContains(self):
r = point.Rect(point.Point(1, 1), point.Point(3, 3))
self.assertTrue(r.contains_point(point.Point(2, 2)))
self.assertFalse(r.contains_circle(point.Point(2, 2), 5))
self.assertFalse(r.contains_point(point.Point(4, 4)))
self.assertFalse(r.contains_circle(point.Point(4, 4), 5))
def testArea(self):
r = point.Rect(point.Point(1, 1), point.Point(3, 4))
self.assertEqual(r.area, 6)
