-
Notifications
You must be signed in to change notification settings - Fork 34
/
common.py
120 lines (93 loc) · 3.1 KB
/
common.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
import time
import os
import re
import requests
import logging
from requests.compat import urljoin
op = os.path
data_uri = 'http://data.pyrocko.org/testing/kite/'
data_dir = op.join(op.dirname(op.abspath(__file__)), 'data')
logger = logging.getLogger('kite.testing')
class DownloadError(Exception):
pass
def get_test_data(fn):
def _dir_content(url):
r = requests.get(url)
r.raise_for_status()
entries = re.findall(r'href="([a-zA-Z0-9_.-]+/?)"', r.text)
files = sorted(set(fn for fn in entries
if not fn.endswith('/')))
return files
def _file_size(url):
r = requests.head(url, headers={'Accept-Encoding': 'identity'})
r.raise_for_status()
print r.headers
return int(r.headers['content-length'])
def _download_file(url, fn_local):
print url, fn_local
if op.exists(fn_local):
if os.stat(fn_local).st_size == _file_size(url):
return fn_local
logger.info('Downloading %s...' % url)
fsize = _file_size(url)
r = requests.get(url, stream=True)
r.raise_for_status()
dl_bytes = 0
with open(fn_local, 'w') as f:
for d in r.iter_content(chunk_size=1024):
dl_bytes += len(d)
f.write(d)
if dl_bytes != fsize:
print url, dl_bytes, fsize
raise DownloadError('Download incomplete!')
logger.info('Download completed!')
return fn_local
url = urljoin(data_uri, fn)
dl_dir = data_dir
if fn.endswith('/'):
dl_dir = op.join(data_dir, fn)
dl_files = _dir_content(url)
dl_files = zip([urljoin(url, u) for u in dl_files],
[op.join(dl_dir, f) for f in dl_files])
else:
dl_files = (url, fn)
return _download_file(*dl_files)
try:
os.makedirs(dl_dir)
except OSError:
pass
return [_download_file(*f) for f in dl_files]
class Benchmark(object):
def __init__(self, prefix=None):
self.prefix = prefix or ''
self.results = []
def __call__(self, func):
def stopwatch(*args):
t0 = time.time()
name = self.prefix + func.__name__
result = func(*args)
elapsed = time.time() - t0
self.results.append((name, elapsed))
return result
return stopwatch
def __str__(self):
rstr = ['Benchmark results']
if self.prefix != '':
rstr[-1] += ' - %s' % self.prefix
if len(self.results) > 0:
indent = max([len(name) for name, _ in self.results])
else:
indent = 0
rstr.append('=' * (indent + 17))
rstr.insert(0, rstr[-1])
for res in self.results:
rstr.append(
'{0:<{indent}}{1:.8f} s'.format(*res, indent=indent+5))
if len(self.results) == 0:
rstr.append('None ran!')
return '\n'.join(rstr)
def setLogLevel(level):
print('SETTING LOGLEVEL!')
level = getattr(logging, level, 'DEBUG')
logging.basicConfig(
level=level)