1# -*- coding: utf-8 -*-
2
3import os
4import time
5
6import tables
7
8tref = time.time()
9trel = tref
10
11
12def show_mem(explain):
13    global tref, trel
14
15    filename = "/proc/%s/status" % os.getpid()
16    with open(filename) as fd:
17        for line in fd:
18            if line.startswith("VmSize:"):
19                vmsize = int(line.split()[1])
20            elif line.startswith("VmRSS:"):
21                vmrss = int(line.split()[1])
22            elif line.startswith("VmData:"):
23                vmdata = int(line.split()[1])
24            elif line.startswith("VmStk:"):
25                vmstk = int(line.split()[1])
26            elif line.startswith("VmExe:"):
27                vmexe = int(line.split()[1])
28            elif line.startswith("VmLib:"):
29                vmlib = int(line.split()[1])
30
31    print("\nMemory usage: ******* %s *******" % explain)
32    print("VmSize: %7s kB\tVmRSS: %7s kB" % (vmsize, vmrss))
33    print("VmData: %7s kB\tVmStk: %7s kB" % (vmdata, vmstk))
34    print("VmExe:  %7s kB\tVmLib: %7s kB" % (vmexe, vmlib))
35    print("WallClock time:", time.time() - tref, end=' ')
36    print("  Delta time:", time.time() - trel)
37    trel = time.time()
38
39
40def write_group(filename, nchildren, niter):
41    for i in range(niter):
42        fileh = tables.open_file(filename, mode="w")
43        for child in range(nchildren):
44            fileh.create_group(fileh.root, 'group' + str(child),
45                               "child: %d" % child)
46        show_mem("After creating. Iter %s" % i)
47        fileh.close()
48        show_mem("After close")
49
50
51def read_group(filename, nchildren, niter):
52    for i in range(niter):
53        fileh = tables.open_file(filename, mode="r")
54        for child in range(nchildren):
55            node = fileh.get_node(fileh.root, 'group' + str(child))
56            assert node is not None
57            # flavor = node._v_attrs.CLASS
58#         for child in fileh.walk_nodes():
59#             pass
60        show_mem("After reading metadata. Iter %s" % i)
61        fileh.close()
62        show_mem("After close")
63
64
65def write_array(filename, nchildren, niter):
66    for i in range(niter):
67        fileh = tables.open_file(filename, mode="w")
68        for child in range(nchildren):
69            fileh.create_array(fileh.root, 'array' + str(child),
70                               [1, 1], "child: %d" % child)
71        show_mem("After creating. Iter %s" % i)
72        fileh.close()
73        show_mem("After close")
74
75
76def read_array(filename, nchildren, niter):
77    for i in range(niter):
78        fileh = tables.open_file(filename, mode="r")
79        for child in range(nchildren):
80            node = fileh.get_node(fileh.root, 'array' + str(child))
81            # flavor = node._v_attrs.FLAVOR
82            data = node[:]  # Read data
83            assert data is not None
84        show_mem("After reading data. Iter %s" % i)
85#         for child in range(nchildren):
86#             node = fileh.get_node(fileh.root, 'array' + str(child))
87#             flavor = node._v_attrs.FLAVOR
88            # flavor = node._v_attrs
89#         for child in fileh.walk_nodes():
90#             pass
91#         show_mem("After reading metadata. Iter %s" % i)
92        fileh.close()
93        show_mem("After close")
94
95
96def write_carray(filename, nchildren, niter):
97    for i in range(niter):
98        fileh = tables.open_file(filename, mode="w")
99        for child in range(nchildren):
100            fileh.create_carray(fileh.root, 'array' + str(child),
101                                tables.IntAtom(), (2,), "child: %d" % child)
102        show_mem("After creating. Iter %s" % i)
103        fileh.close()
104        show_mem("After close")
105
106
107def read_carray(filename, nchildren, niter):
108    for i in range(niter):
109        fileh = tables.open_file(filename, mode="r")
110        for child in range(nchildren):
111            node = fileh.get_node(fileh.root, 'array' + str(child))
112            # flavor = node._v_attrs.FLAVOR
113            data = node[:]  # Read data
114            assert data is not None
115            # print("data-->", data)
116        show_mem("After reading data. Iter %s" % i)
117        fileh.close()
118        show_mem("After close")
119
120
121def write_earray(filename, nchildren, niter):
122    for i in range(niter):
123        fileh = tables.open_file(filename, mode="w")
124        for child in range(nchildren):
125            ea = fileh.create_earray(fileh.root, 'array' + str(child),
126                                     tables.IntAtom(), shape=(0,),
127                                     title="child: %d" % child)
128            ea.append([1, 2, 3])
129        show_mem("After creating. Iter %s" % i)
130        fileh.close()
131        show_mem("After close")
132
133
134def read_earray(filename, nchildren, niter):
135    for i in range(niter):
136        fileh = tables.open_file(filename, mode="r")
137        for child in range(nchildren):
138            node = fileh.get_node(fileh.root, 'array' + str(child))
139            # flavor = node._v_attrs.FLAVOR
140            data = node[:]  # Read data
141            assert data is not None
142            # print("data-->", data)
143        show_mem("After reading data. Iter %s" % i)
144        fileh.close()
145        show_mem("After close")
146
147
148def write_vlarray(filename, nchildren, niter):
149    for i in range(niter):
150        fileh = tables.open_file(filename, mode="w")
151        for child in range(nchildren):
152            vl = fileh.create_vlarray(fileh.root, 'array' + str(child),
153                                      tables.IntAtom(), "child: %d" % child)
154            vl.append([1, 2, 3])
155        show_mem("After creating. Iter %s" % i)
156        fileh.close()
157        show_mem("After close")
158
159
160def read_vlarray(filename, nchildren, niter):
161    for i in range(niter):
162        fileh = tables.open_file(filename, mode="r")
163        for child in range(nchildren):
164            node = fileh.get_node(fileh.root, 'array' + str(child))
165            # flavor = node._v_attrs.FLAVOR
166            data = node[:]  # Read data
167            assert data is not None
168            # print("data-->", data)
169        show_mem("After reading data. Iter %s" % i)
170        fileh.close()
171        show_mem("After close")
172
173
174def write_table(filename, nchildren, niter):
175
176    class Record(tables.IsDescription):
177        var1 = tables.IntCol(pos=1)
178        var2 = tables.StringCol(length=1, pos=2)
179        var3 = tables.FloatCol(pos=3)
180
181    for i in range(niter):
182        fileh = tables.open_file(filename, mode="w")
183        for child in range(nchildren):
184            t = fileh.create_table(fileh.root, 'table' + str(child),
185                                   Record, "child: %d" % child)
186            t.append([[1, "2", 3.]])
187        show_mem("After creating. Iter %s" % i)
188        fileh.close()
189        show_mem("After close")
190
191
192def read_table(filename, nchildren, niter):
193    for i in range(niter):
194        fileh = tables.open_file(filename, mode="r")
195        for child in range(nchildren):
196            node = fileh.get_node(fileh.root, 'table' + str(child))
197            # klass = node._v_attrs.CLASS
198            data = node[:]  # Read data
199            assert data is not None
200            # print("data-->", data)
201        show_mem("After reading data. Iter %s" % i)
202        fileh.close()
203        show_mem("After close")
204
205
206def write_xtable(filename, nchildren, niter):
207
208    class Record(tables.IsDescription):
209        var1 = tables.IntCol(pos=1)
210        var2 = tables.StringCol(length=1, pos=2)
211        var3 = tables.FloatCol(pos=3)
212
213    for i in range(niter):
214        fileh = tables.open_file(filename, mode="w")
215        for child in range(nchildren):
216            t = fileh.create_table(fileh.root, 'table' + str(child),
217                                   Record, "child: %d" % child)
218            t.append([[1, "2", 3.]])
219            t.cols.var1.create_index()
220        show_mem("After creating. Iter %s" % i)
221        fileh.close()
222        show_mem("After close")
223
224
225def read_xtable(filename, nchildren, niter):
226    for i in range(niter):
227        fileh = tables.open_file(filename, mode="r")
228        for child in range(nchildren):
229            node = fileh.get_node(fileh.root, 'table' + str(child))
230            # klass = node._v_attrs.CLASS
231            # data = node[:]  # Read data
232            # print("data-->", data)
233        show_mem("After reading data. Iter %s" % i)
234        fileh.close()
235        show_mem("After close")
236        del node
237
238
239if __name__ == '__main__':
240    import pstats
241    import argparse
242    import profile as prof
243
244    def _get_parser():
245        parser = argparse.ArgumentParser(
246            description='Check for PyTables memory leaks.')
247        parser.add_argument('-v', '--verbose', action='store_true',
248                            help='enable verbose mode')
249        parser.add_argument('-p', '--profile', action='store_true',
250                            help='profile')
251        parser.add_argument('-a', '--array', action='store_true',
252                            help='create/read arrays (default)')
253        parser.add_argument('-c', '--carray', action='store_true',
254                            help='create/read carrays')
255        parser.add_argument('-e', '--earray', action='store_true',
256                            help='create/read earrays')
257        parser.add_argument('-l', '--vlarray', action='store_true',
258                            help='create/read vlarrays')
259        parser.add_argument('-t', '--table', action='store_true',
260                            help='create/read tables')
261        parser.add_argument('-x', '--indexed-table', action='store_true',
262                            dest='xtable', help='create/read indexed-tables')
263        parser.add_argument('-g', '--group', action='store_true',
264                            help='create/read groups')
265        parser.add_argument('-r', '--read', action='store_true',
266                            help='only read test')
267        parser.add_argument('-w', '--write', action='store_true',
268                            help='only write test')
269        parser.add_argument('-n', '--nchildren', type=int, default=1000,
270                            help='number of children (%(default)d is the '
271                                 'default)')
272        parser.add_argument('-i', '--niter', type=int, default=3,
273                            help='number of iterations (default: %(default)d)')
274
275        parser.add_argument('filename', help='HDF5 file name')
276
277        return parser
278
279    parser = _get_parser()
280    args = parser.parse_args()
281
282    # set 'array' as default value if no ather option has been specified
283    for name in ('carray', 'earray', 'vlarray', 'table', 'xtable', 'group'):
284        if getattr(args, name):
285            break
286    else:
287        args.array = True
288
289    filename = args.filename
290    nchildren = args.nchildren
291    niter = args.niter
292
293    if args.array:
294        fwrite = 'write_array'
295        fread = 'read_array'
296    elif args.carray:
297        fwrite = 'write_carray'
298        fread = 'read_carray'
299    elif args.earray:
300        fwrite = 'write_earray'
301        fread = 'read_earray'
302    elif args.vlarray:
303        fwrite = 'write_vlarray'
304        fread = 'read_vlarray'
305    elif args.table:
306        fwrite = 'write_table'
307        fread = 'read_table'
308    elif args.xtable:
309        fwrite = 'write_xtable'
310        fread = 'read_xtable'
311    elif args.group:
312        fwrite = 'write_group'
313        fread = 'read_group'
314
315    show_mem("Before open")
316    if args.write:
317        if args.profile:
318            prof.run(str(fwrite)+'(filename, nchildren, niter)',
319                     'write_file.prof')
320            stats = pstats.Stats('write_file.prof')
321            stats.strip_dirs()
322            stats.sort_stats('time', 'calls')
323            if args.verbose:
324                stats.print_stats()
325            else:
326                stats.print_stats(20)
327        else:
328            eval(fwrite+'(filename, nchildren, niter)')
329    if args.read:
330        if args.profile:
331            prof.run(fread+'(filename, nchildren, niter)', 'read_file.prof')
332            stats = pstats.Stats('read_file.prof')
333            stats.strip_dirs()
334            stats.sort_stats('time', 'calls')
335            if args.verbose:
336                print('profile -verbose')
337                stats.print_stats()
338            else:
339                stats.print_stats(20)
340        else:
341            eval(fread+'(filename, nchildren, niter)')
342