utils.py 文件源码

def load_bundle(filename, regex="^\d+", chunk=None):
    """ Load graphs from zip archive """

    pool = Pool()
    archive = zf.ZipFile(filename)

    # Determine entries and select subset if requested
    entries = archive.namelist()
    if chunk:
        entries = list(set(entries) & set(chunk))
    entries = [(archive, entry) for entry in entries]

    # Load entries in parallel
    func = partial(load_bundle_entry, regex=re.compile(regex))
    items = pool.map(func, entries)
    items = filter(lambda (g, l): g is not None, items)
    graphs, labels = zip(*items)

    archive.close()
    pool.close()
    pool.join()

    return graphs, labels