def __init__(self, world, root=None, subtrees=None, data_size=0):
"""Initialize a new PopulationTree.
If <world> is True, then this tree is the root of the population tree,
and it should load data from the World Bank API.
In this case, none of the other parameters are used.
If <world> is False, pass the other arguments directly to the superclass
constructor. Do NOT load new data from the World Bank API.
@type self: PopulationTree
@type world: bool
@type root: object
@type subtrees: list[PopulationTree] | None
@type data_size: int
>>> t1 = PopulationTree(True)
>>> t1._root
'world'
>>> t2 = PopulationTree(False, 'North America')
>>> t2._root
'North America'
"""
if world:
region_trees = _load_data()
AbstractTree.__init__(self, 'World', region_trees)
else:
if subtrees is None:
subtrees = []
AbstractTree.__init__(self, root, subtrees, data_size)
def event_loop(screen: pygame.Surface, tree: AbstractTree) -> None:
"""Respond to events (mouse clicks, key presses) and update the display.
Note that the event loop is an *infinite loop*: it continually waits for
the next event, determines the event's type, and then updates the state
of the visualisation or the tree itself, updating the display if necessary.
This loop ends when the user closes the window.
"""
# We strongly recommend using a variable to keep track of the currently-
# selected leaf (type AbstractTree | None).
# But feel free to remove it, and/or add new variables, to help keep
# track of the state of the program.
selected_leaf = None
text = ''
while True:
# Wait for an event
event = pygame.event.poll()
if event.type == pygame.QUIT:
return
if event.type == pygame.MOUSEBUTTONUP and event.button == 1: # left click to select or deselect
leaf = tree.find_rect(pygame.mouse.get_pos())
if selected_leaf == leaf: # deselect
selected_leaf = None
text = ''
render_display(screen, tree, text)
else: # select
selected_leaf = leaf
text = selected_leaf.get_path() + ' ' + '(' + \
str(selected_leaf.data_size) + ')'
render_display(screen, tree, text)
if event.type == pygame.MOUSEBUTTONUP and event.button == 3: # right click to delete
leaf = tree.find_rect(pygame.mouse.get_pos())
leaf.del_update_parents()
leaf.delete_leaf()
render_display(screen, tree, text)
if event.type == pygame.KEYUP: # Key Press
if selected_leaf:
if event.key == pygame.K_UP: # Up Arrow
selected_leaf.adjust_size(True)
text = selected_leaf.get_path() + ' ' + '(' + \
str(selected_leaf.data_size) + ')'
render_display(screen, tree, text)
if event.key == pygame.K_DOWN: # Down Arrow
selected_leaf.adjust_size(False)
text = selected_leaf.get_path() + ' ' + '(' + \
str(selected_leaf.data_size) + ')'
render_display(screen, tree, text)
def __init__(self, path):
if path is None:
self._root, self._subtrees = None
AbstractTree.__init__(self, root=None, subtrees=[])
new_path = path.split('/')
if len(new_path) == 1:
self._root = os.path.basename(path)
self._subtrees = []
else:
self._root = os.path.basename(path)
self._subtrees = new_path[2:]
total_data = 0
for tree in self._subtrees:
tree_size = os.path.getsize(tree)
tree.data_size = tree_size
total_data += tree_size
self.data_size = total_data
def __init__(self: PopulationTree,
world: bool,
root: Optional[object] = None,
subtrees: Optional[List[PopulationTree]] = None,
data_size: int = 0) -> None:
"""Initialize a new PopulationTree.
If <world> is True, then this tree is the root of the population tree,
and it should load data from the World Bank files.
In this case, none of the other parameters are used.
If <world> is False, pass the other arguments directly to the superclass
constructor. Do NOT load new data from the World Bank files.
"""
if world:
region_trees = _load_data()
AbstractTree.__init__(self, 'World', region_trees)
else:
if subtrees is None:
subtrees = []
AbstractTree.__init__(self, root, subtrees, data_size)
def __init__(self, world, root=None, subtrees=None, data_size=0):
"""Initialize a new PopulationTree.
If <world> is True, then this tree is the root of the population tree,
and it should load data from the World Bank API.
If <world> is False, pass the other arguments directly to the superclass
constructor.
@type self: PopulationTree
@type world: bool
@type root: object
@type subtrees: list[PopulationTree] | None
@type data_size: int
"""
if world:
region_trees = _load_data()
AbstractTree.__init__(self, 'World', region_trees)
else:
if subtrees is None:
subtrees = []
AbstractTree.__init__(self, root, subtrees, data_size)
def render_display(screen: pygame.Surface, tree: AbstractTree,
text: str) -> None:
"""Render a treemap and text display to the given screen.
Use the constants TREEMAP_HEIGHT and FONT_HEIGHT to divide the
screen vertically into the treemap and text comments.
"""
# First, clear the screen
pygame.draw.rect(screen, pygame.color.THECOLORS['black'],
(0, 0, WIDTH, HEIGHT))
# TODO: Implement this function!
rectangles = tree.generate_treemap((0, 0, WIDTH, TREEMAP_HEIGHT))
for rectangle in rectangles:
pygame.draw.rect(screen, rectangle[1], rectangle[0])
_render_text(screen, text)
# This must be called *after* all other pygame functions have run.
pygame.display.flip()