-
Notifications
You must be signed in to change notification settings - Fork 0
/
tm_trees.py
377 lines (313 loc) · 12.6 KB
/
tm_trees.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
"""Trees for Treemap
=== Module Description ===
This module contains the basic tree interface required by the treemap
visualiser. You will both add to the abstract class, and complete a
concrete implementation of a subclass to represent files and folders on your
computer's file system.
"""
from __future__ import annotations
import os
import math
from random import randint
from typing import List, Tuple, Optional, Union
WIDTH = 800
HEIGHT = 600
class TMTree:
"""A TreeMappableTree: a tree that is compatible with the treemap
visualiser.
This is an abstract class that should not be instantiated directly.
You may NOT add any attributes, public or private, to this class.
However, part of this assignment will involve you implementing new public
*methods* for this interface.
You should not add any new public methods other than those required by
the client code.
You can, however, freely add private methods as needed.
=== Public Attributes ===
rect:
The pygame rectangle representing this node in the treemap
visualization.
data_size:
The size of the data represented by this tree.
=== Private Attributes ===
_colour:
The RGB colour value of the root of this tree.
__name__:
The root value of this tree, or None if this tree is empty.
_subtrees:
The subtrees of this tree.
_parent_tree:
The parent tree of this tree; i.e., the tree that contains this tree
as a subtree, or None if this tree is not part of a larger tree.
_expanded:
Whether or not this tree is considered expanded for visualization.
=== Representation Invariants ===
- data_size >= 0
- If _subtrees is not empty, then data_size is equal to the sum of the
data_size of each subtree.
- _colour's elements are each in the range 0-255.
- If _name is None, then _subtrees is empty, _parent_tree is None, and
data_size is 0.
This setting of attributes represents an empty tree.
- if _parent_tree is not None, then self is in _parent_tree._subtrees
- if _expanded is True, then _parent_tree._expanded is True
- if _expanded is False, then _expanded is False for every tree
in _subtrees
- if _subtrees is empty, then _expanded is False
"""
rect: Tuple[int, int, int, int]
data_size: int
_colour: Tuple[int, int, int]
_name: str
_subtrees: List[TMTree]
_parent_tree: Optional[TMTree]
_expanded: bool
def __init__(self, name: str, subtrees: List[TMTree],
data_size: int = 0) -> None:
"""Initialize a new TMTree with a random colour and the provided <name>.
If <subtrees> is empty, use <data_size> to initialize this tree's
data_size.
If <subtrees> is not empty, ignore the parameter <data_size>,
and calculate this tree's data_size instead.
Set this tree as the parent for each of its subtrees.
Precondition: if <name> is None, then <subtrees> is empty.
"""
self.rect = (0, 0, 0, 0)
self._name = name
self._subtrees = subtrees[:]
self._parent_tree = None
self._expanded = False
self._colour = (randint(0, 255), randint(0, 255), randint(0, 255))
if self._is_leaf():
self.data_size = data_size
else:
self.data_size = 0
for subtree in self._subtrees:
self.data_size += subtree.data_size
subtree._parent_tree = self
def is_empty(self) -> bool:
"""Return True iff this tree is empty.
"""
return self._name is None
def update_rectangles(self, rect: Tuple[int, int, int, int]) -> None:
"""Update the rectangles in this tree and its descendents using the
treemap algorithm to fill the area defined by pygame rectangle <rect>.
"""
if self.data_size == 0:
self.rect = (0, 0, 0, 0)
else:
self.rect = rect
if not self._is_leaf():
x, y, w, h = rect;
max_x, max_y = w, h
if w > h:
a = 0
for subtree in self._subtrees[:-1]:
ratio = self._check(subtree)
subtree.update_rectangles((x + a, y,
math.floor(ratio * w), h))
a += subtree.rect[2]
self._subtrees[-1].update_rectangles((x + a, y,
max_x - a, max_y))
else:
b = 0
for subtree in self._subtrees[:-1]:
ratio = self._check(subtree)
subtree.update_rectangles((x, y + b,
w, math.floor(ratio * h)))
b += subtree.rect[3]
self._subtrees[-1].update_rectangles((x, y + b,
max_x, max_y - b))
def _check(self, subtree: TMTree) -> float:
"""Check if self size is 0
"""
if self.data_size == 0:
return 0.0
return subtree.data_size / self.data_size
def get_rectangles(self) -> List[Tuple[Tuple[int, int, int, int],
Tuple[int, int, int]]]:
"""Return a list with tuples for every leaf in the displayed-tree
rooted at this tree. Each tuple consists of a tuple that defines the
appropriate pygame rectangle to display for a leaf, and the colour
to fill it with.
"""
if self.data_size == 0:
return []
if not self._expanded:
return [(self.rect, self._colour)]
else:
if self._is_leaf():
return [(self.rect, self._colour)]
r = []
for subtree in self._subtrees:
r.extend(subtree.get_rectangles())
return r
def get_tree_at_position(self, pos: Tuple[int, int]) -> Optional[TMTree]:
"""Return the leaf in the displayed-tree rooted at this tree whose
rectangle contains position <pos>, or None if <pos> is outside of this
tree's rectangle.
If <pos> is on the shared edge between two rectangles, return the
tree represented by the rectangle that is closer to the origin.
"""
if pos > (WIDTH, HEIGHT) or self.is_empty() or not self._includes(pos):
return None
if not (self._expanded and self._subtrees):
return self
result = None
for sub in self._subtrees:
x = sub.get_tree_at_position(pos)
if x:
if not result or ((x.rect[0] < result.rect[0]) or
(x.rect[1] < result.rect[1])):
result = x
return result
def _includes(self, pos: Tuple[int, int]) -> bool:
return (self.rect[0] <= pos[0] <= (self.rect[0] + self.rect[2]) and
self.rect[1] <= pos[1] <= (self.rect[1] + self.rect[3]))
def update_data_sizes(self) -> int:
"""Update the data_size for this tree and its subtrees, based on the
size of their leaves, and return the new size.
If this tree is a leaf, return its size unchanged.
"""
if self._is_leaf():
return self.data_size
r = 0
for subtree in self._subtrees:
r += subtree.update_data_sizes()
self.data_size = r
return self.data_size
def move(self, destination: TMTree) -> None:
"""If this tree is a leaf, and <destination> is not a leaf, move this
tree to be the last subtree of <destination>. Otherwise, do nothing.
"""
if self._movable(destination):
item = self
if self._parent_tree:
self._parent_tree._subtrees.remove(item)
if not self._parent_tree._subtrees:
self._parent_tree.data_size = 0
self._parent_tree.update_data_sizes()
destination._subtrees.append(item)
item._parent_tree = destination
destination.update_data_sizes()
def change_size(self, factor: float) -> None:
"""Change the value of this tree's data_size attribute by <factor>.
Always round up the amount to change, so that it's an int, and
some change is made.
Do nothing if this tree is not a leaf.
"""
if self._is_leaf():
change = factor * self.data_size
if change > 0:
self.data_size += math.ceil(change)
else:
if self.data_size + math.floor(change) <= 1:
self.data_size = 1
else:
self.data_size += math.floor(change)
if not self._is_root():
self._parent_tree.update_data_sizes()
def expand(self) -> None:
"""Expand the file."""
if not self._subtrees:
return
self._expanded = True
item = self
while item:
item._expanded = True
item = item._parent_tree
def expand_all(self) -> None:
"""Expand all files."""
if self._subtrees:
self._expanded = True
for sub in self._subtrees:
sub.expand_all()
def collapse(self) -> None:
"""Unexpanded the file.
"""
if self._parent_tree:
self._parent_tree._real_collapse()
def collapse_all(self) -> None:
"""Unexpanded all files.
"""
if self._parent_tree:
r = self
while r._parent_tree:
r = r._parent_tree
r._real_collapse()
def _movable(self, destination: TMTree) -> bool:
return self._is_leaf() and not destination._is_leaf()
def _real_collapse(self) -> None:
"""Helper for collapse.
"""
self._expanded = False
for sub in self._subtrees:
sub._real_collapse()
def _is_root(self) -> bool:
return not self._parent_tree
def _is_leaf(self) -> bool:
return not self._subtrees
# Methods for the string representation
def get_path_string(self, final_node: bool = True) -> str:
"""Return a string representing the path containing this tree
and its ancestors, using the separator for this tree between each
tree's name. If <final_node>, then add the suffix for the tree.
"""
if self._parent_tree is None:
path_str = self._name
if final_node:
path_str += self.get_suffix()
return path_str
else:
path_str = (self._parent_tree.get_path_string(False) +
self.get_separator() + self._name)
if final_node or len(self._subtrees) == 0:
path_str += self.get_suffix()
return path_str
def get_separator(self) -> str:
"""Return the string used to separate names in the string
representation of a path from the tree root to this tree.
"""
raise NotImplementedError
def get_suffix(self) -> str:
"""Return the string used at the end of the string representation of
a path from the tree root to this tree.
"""
raise NotImplementedError
class FileSystemTree(TMTree):
"""A tree representation of files and folders in a file system.
The internal nodes represent folders, and the leaves represent regular
files (e.g., PDF documents, movie files, Python source code files, etc.).
The _name attribute stores the *name* of the folder or file, not its full
path. E.g., store 'Documents', not '/Users/John/Documents'
The data_size attribute for regular files is simply the size of the file,
as reported by os.path.getsize.
"""
def __init__(self, path: str) -> None:
"""Store the file tree structure contained in the given file or folder.
Precondition: <path> is a valid path for this computer.
"""
subtrees = []
name = os.path.basename(path)
if os.path.isdir(path):
dirs = os.listdir(path)
for file in dirs:
sub_path = os.path.join(path, file)
subtree = FileSystemTree(sub_path)
subtrees.append(subtree)
TMTree.__init__(self, name, subtrees, os.path.getsize(path))
def get_separator(self) -> str:
"""Return the file separator for this OS.
"""
return os.sep
def get_suffix(self) -> str:
"""Return the final descriptor of this tree.
"""
if len(self._subtrees) == 0:
return ' (file)'
else:
return ' (folder)'
if __name__ == '__main__':
pass
# import doctest
#
# doctest.testmod()