# bk.py
"""
Bookkeeping helper functions.
"""
try:
import cPickle as pickle
except:
import pickle
try:
from itertools import izip
except:
izip = zip
from glob import glob
import numpy as np
import os
import numpy as np
import sys
def load(filename):
return np.load(filename)
#with open(filename) as f:
# return pickle.load(f)
[docs]
def get_all(root_dir, pattern, merge=np.sum,
pre_convert=None,
post_convert=None,
log=True):
"""
Get a tree of nested dictionaries with files loaded as leaf items.
Args:
root_dir (str): the root of the directory tree
pattern (str): a glob pattern for the files to load
merge (func(list) -> object): function that merges loaded objects
"""
pre_convert = pre_convert or (lambda x: x)
post_convert = post_convert or (lambda x: x)
def note_file(path):
s = '{} ...'.format(path)
note_file.maxlen = max(note_file.maxlen, len(s))
note_file.n_loaded += 1
if log:
print('\r'+s,)
sys.stdout.flush()
note_file.maxlen = 0
note_file.n_loaded = 0
def load(filename):
note_file(filename)
return np.load(filename)
def type_ok(typename, values):
if typename == bool:
for value in values:
if value not in ('True', 'False'):
return False
return True
else:
try:
list(map(typename, values))
return True
except:
return False
def get(dir, pattern, merge):
"""
Recursively descend into dir, adding an item for merged objects
matching pattern, or otherwise adding dicts for each subdir.
"""
paths = [d
for d in sorted(glob(os.path.join(dir, '*')))
if os.path.isdir(d)]
if paths:
out = {}
basenames = list(map(os.path.basename, paths))
for typename in (bool, int, float, str):
if type_ok(typename, basenames):
break
keys = list(map(typename, basenames))
for (key, path) in izip(keys, paths):
result = get(path, pattern, merge)
if isinstance(result, np.ndarray) or result:
out[key] = result
return out
else:
paths = sorted(glob(os.path.join(dir, pattern)))
load_and_convert = lambda f: pre_convert(load(f))
objects = list(map(load_and_convert, paths))
if objects:
return post_convert(merge(objects))
else:
return False
out = get(root_dir, pattern, merge)
if log:
print('\r'+note_file.maxlen * ' ',)
print('\r{} files loaded.'.format(note_file.n_loaded))
return out
[docs]
def get_best(d, *values):
"""
Descend into dict-tree d, obtaining the best-match subtree or item.
"""
value = values[0]
if value in (np.inf, -np.inf):
key = value
else:
keys = np.sort(list(d.keys()))
idx = np.argwhere(keys == value)
if idx.size or (value in keys):
key = keys[idx[0][0]]
else:
key = keys[np.argmin(np.abs(value - keys))]
if len(values) > 1:
return get_best(d[key], *values[1:])
else:
return d[key]