-
Notifications
You must be signed in to change notification settings - Fork 0
/
Memory_Tree.py
175 lines (164 loc) · 4.43 KB
/
Memory_Tree.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
#!/usr/bin/python3
import os
from itertools import chain
import textwrap
import time
import heapq
import logging
logging.basicConfig(filename='log.log',level=logging.DEBUG)
class node(object):
def __init__(self,name):
self.memsize=0
self.name=name
self.kiddir=[]
self.kidfile=[]
self.dist=0
self.depth=0
def add_subdir(self,kid_node):
self.kiddir.append(kid_node)
def add_subfile(self,kid_node):
self.kidfile.append(kid_node)
def update_memsize(self,memory):
self.memsize=memory
def update_dist(self,dist):
self.dist=dist
def update_depth(self,depth):
self.depth=depth
class Mem_Tree(object):
def __init__(self):
self.root=None
self.depth=0
self._list=[]
def initialize(self,rootpath):
if len(rootpath)==1: #If only one rootpath provided
Root=node(rootpath[0])
self.root=Root
waitlist=[Root]
else:
Root=node('&&'.join(rootpath)) #If several path are provided in a list
self.root=Root
waitlist=[]
for i in rootpath:
new_node=node(i)
if os.path.isdir(i):
Root.add_subdir(new_node)
else:
Root.add_subfile(new_node)
waitlist.append(new_node)
while waitlist:
cur_node=waitlist.pop()
cur_path=cur_node.name
try:
names=os.listdir(cur_path)
except PermissionError as err:
logging.warning(err)
for name in names:
complete_path=os.path.join(cur_path,name)
new_node=node(complete_path)
if os.path.isdir(complete_path):
cur_node.add_subdir(new_node)
waitlist.append(new_node)
else:
cur_node.add_subfile(new_node)
try:
memory=os.path.getsize(complete_path)
new_node.update_memsize(memory)
tmp={}
tmp['name']=new_node.name
tmp['mem']=memory
self._list.append(tmp)
except PermissionError as err:
logging.warning(err)
except FileNotFoundError as err:
logging.warning(err)
def search(self,filename):
#Search the path for a certain filename
DFS=[self.root]
result=[]
while DFS:
cur_node=DFS.pop()
if os.path.isdir(cur_node.name):
try_path=os.path.join(cur_node.name,filename)
if os.path.exists(try_path):
result.append(try_path)
for k in cur_node.kiddir:
DFS.append(k)
if not result:
print("Sorry, mission failed. T_T")
else:
print("The file (%s) is found in the following place(s)" % filename)
for i in result:
print(i)
def update_memory_all(self):
#Update the memory used in each node
DFS=[self.root]
DFS1=[]
dist=0
depth=0
while DFS:
cur_node=DFS.pop()
DFS1.append(cur_node)
for k in chain(cur_node.kiddir,cur_node.kidfile):
DFS.append(k)
k.update_dist(cur_node.dist+1)
while DFS1:
n=DFS1.pop()
if os.path.isdir(n.name):
result=0
depth=0
for k in chain(n.kiddir,n.kidfile):
result+=k.memsize
if k.depth>depth:
depth=k.depth
depth+=1
n.update_memsize(result)
n.update_depth(depth)
self.depth=self.root.depth
def maxN(self,N):
#Find the max N memory consuming files
top=heapq.nlargest(N,self._list,key=lambda s:s['mem'])
print("The max (%d) memory consuming files are: " %(N))
for i in top:
print(i)
def pretty_print(self,m=True,f=False):
#Print the files and folders under root, with their memory usage if m(with_memory)=True
#Only print folders is f(folders_only)=True
base=3
width=self.depth*base+50
DFS=[self.root]
while DFS:
cur_node=DFS.pop()
indent_size=cur_node.dist*base
name=os.path.basename(cur_node.name)
star_size=width-indent_size-len(name)
try:
if m:
print(textwrap.indent(name,''.join([' ']*indent_size)),''.join(['.']*star_size),cur_node.memsize)
else:
print(textwrap.indent(name,''.join([' ']*indent_size)))
except UnicodeEncodeError as err:
logging.debug(err)
for k in cur_node.kiddir:
DFS.append(k)
if not f:
for k in cur_node.kidfile:
indent_size=k.dist*base
name=os.path.basename(k.name)
star_size=width-indent_size-len(name)
try:
if m:
print(textwrap.indent(name,''.join([' ']*indent_size)),''.join(['.']*star_size),k.memsize)
else:
print(textwrap.indent(name,''.join([' ']*indent_size)))
except UnicodeEncodeError as err:
print ('Handling unicode encode error',err)
# start_time=time.time()
# path=['D:/','C:/']
# T=Tree()
# T.initialize(path)
# T.search('UG_PA13130.pdf')
# T.maxN(10)
# T.update_memory_all()
# T.pretty_print(folders_only=True)
# end_time=time.time()
# print("Running time is: ",end_time-start_time)