def test_save_zip(tmpdir):
# load from zipped SBT, save to zipped SBT, and then search.
testdata = utils.get_test_data("v6.sbt.zip")
testsbt = tmpdir.join("v6.sbt.zip")
newsbt = tmpdir.join("new.sbt.zip")
shutil.copyfile(testdata, str(testsbt))
tree = SBT.load(str(testsbt), leaf_loader=SigLeaf.load)
tree.save(str(newsbt))
assert newsbt.exists()
new_tree = SBT.load(str(newsbt), leaf_loader=SigLeaf.load)
assert isinstance(new_tree.storage, ZipStorage)
assert new_tree.storage.list_sbts() == ['new.sbt.json']
to_search = load_one_signature(utils.get_test_data(utils.SIG_FILES[0]))
print("*" * 60)
print("{}:".format(to_search))
old_result = {str(s) for s in tree.find(search_minhashes, to_search, 0.1)}
new_result = {
str(s)
for s in new_tree.find(search_minhashes, to_search, 0.1)
}
print(*new_result, sep="\n")
assert old_result == new_result
assert len(new_result) == 2
def test_load_future(tmpdir):
with open(str(tmpdir.join("v9999.sbt.json")), 'w') as f:
json.dump({'version': 9999}, f)
with pytest.raises(IndexNotSupported) as excinfo:
SBT.load(str(tmpdir.join("v9999.sbt.json")))
assert "index format is not supported" in str(excinfo.value)
def test_sbt_tarstorage():
factory = GraphFactory(31, 1e5, 4)
with utils.TempDirectory() as location:
tree = SBT(factory)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {str(s) for s in tree.find(search_minhashes,
to_search.data, 0.1)}
print(*old_result, sep='\n')
with TarStorage(os.path.join(location, 'tree.tar.gz')) as storage:
tree.save(os.path.join(location, 'tree'), storage=storage)
with TarStorage(os.path.join(location, 'tree.tar.gz')) as storage:
tree = SBT.load(os.path.join(location, 'tree'),
leaf_loader=SigLeaf.load,
storage=storage)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {str(s) for s in tree.find(search_minhashes,
to_search.data, 0.1)}
print(*new_result, sep='\n')
assert old_result == new_result
def test_tree_save_load(n_children):
factory = GraphFactory(31, 1e5, 4)
tree = SBT(factory, d=n_children)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*old_result, sep='\n')
with utils.TempDirectory() as location:
tree.save(os.path.join(location, 'demo'))
tree = SBT.load(os.path.join(location, 'demo'),
leaf_loader=SigLeaf.load)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*new_result, sep='\n')
assert old_result == new_result
def test_tree_repair():
tree_repair = SBT.load(utils.get_test_data('leaves.sbt.json'),
leaf_loader=SigLeaf.load)
tree_cur = SBT.load(utils.get_test_data('v3.sbt.json'),
leaf_loader=SigLeaf.load)
testdata1 = utils.get_test_data(utils.SIG_FILES[0])
to_search = next(signature.load_signatures(testdata1))
results_repair = {str(s) for s in tree_repair.find(search_minhashes,
to_search, 0.1)}
results_cur = {str(s) for s in tree_cur.find(search_minhashes,
to_search, 0.1)}
assert results_repair == results_cur
assert len(results_repair) == 2
def test_tree_v2_load():
tree_v2 = SBT.load(utils.get_test_data('v2.sbt.json'),
leaf_loader=SigLeaf.load)
tree_cur = SBT.load(utils.get_test_data('v3.sbt.json'),
leaf_loader=SigLeaf.load)
testdata1 = utils.get_test_data(utils.SIG_FILES[0])
to_search = next(signature.load_signatures(testdata1))
results_v2 = {str(s) for s in tree_v2.find(search_minhashes_containment,
to_search, 0.1)}
results_cur = {str(s) for s in tree_cur.find(search_minhashes_containment,
to_search, 0.1)}
assert results_v2 == results_cur
assert len(results_v2) == 4
def test_tree_v2_load():
tree_v2 = SBT.load(utils.get_test_data('v2.sbt.json'),
leaf_loader=SigLeaf.load)
tree_cur = SBT.load(utils.get_test_data('v3.sbt.json'),
leaf_loader=SigLeaf.load)
testdata1 = utils.get_test_data(utils.SIG_FILES[0])
to_search = next(signature.load_signatures(testdata1))
results_v2 = {
str(s)
for s in tree_v2.find(search_minhashes_containment, to_search, 0.1)
}
results_cur = {
str(s)
for s in tree_cur.find(search_minhashes_containment, to_search, 0.1)
}
assert results_v2 == results_cur
assert len(results_v2) == 4
def test_tree_old_load(old_version):
tree_v1 = SBT.load(utils.get_test_data('{}.sbt.json'.format(old_version)),
leaf_loader=SigLeaf.load)
tree_cur = SBT.load(utils.get_test_data('v6.sbt.json'),
leaf_loader=SigLeaf.load)
testdata1 = utils.get_test_data(utils.SIG_FILES[0])
to_search = load_one_signature(testdata1)
results_v1 = {
str(s)
for s in tree_v1.find(search_minhashes_containment, to_search, 0.1)
}
results_cur = {
str(s)
for s in tree_cur.find(search_minhashes_containment, to_search, 0.1)
}
assert results_v1 == results_cur
assert len(results_v1) == 4
def test_tree_repair():
tree_repair = SBT.load(utils.get_test_data('leaves.sbt.json'),
leaf_loader=SigLeaf.load)
tree_cur = SBT.load(utils.get_test_data('v3.sbt.json'),
leaf_loader=SigLeaf.load)
testdata1 = utils.get_test_data(utils.SIG_FILES[0])
to_search = next(signature.load_signatures(testdata1))
results_repair = {
str(s)
for s in tree_repair.find(search_minhashes, to_search, 0.1)
}
results_cur = {
str(s)
for s in tree_cur.find(search_minhashes, to_search, 0.1)
}
assert results_repair == results_cur
assert len(results_repair) == 2
def test_load_zip(tmpdir):
# search zipped SBT
testdata = utils.get_test_data("v6.sbt.zip")
testsbt = tmpdir.join("v6.sbt.zip")
shutil.copyfile(testdata, str(testsbt))
tree = SBT.load(str(testsbt), leaf_loader=SigLeaf.load)
to_search = load_one_signature(utils.get_test_data(utils.SIG_FILES[0]))
print("*" * 60)
print("{}:".format(to_search))
new_result = {str(s) for s in tree.find(search_minhashes, to_search, 0.1)}
print(*new_result, sep="\n")
assert len(new_result) == 2
def test_tree_repair_add_node():
tree_repair = SBT.load(utils.get_test_data('leaves.sbt.json'),
leaf_loader=SigLeaf.load)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree_repair.add_node(leaf)
for pos, node in list(tree_repair.nodes.items()):
# Every parent of a node must be an internal node (and not a leaf),
# except for node 0 (the root), whose parent is None.
if pos != 0:
assert isinstance(tree_repair.parent(pos).node, Node)
# Leaf nodes can't have children
if isinstance(node, Leaf):
assert all(c.node is None for c in tree_repair.children(pos))
def test_tree_repair_add_node():
tree_repair = SBT.load(utils.get_test_data('leaves.sbt.json'),
leaf_loader=SigLeaf.load)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree_repair.add_node(leaf)
for pos, node in list(tree_repair.nodes.items()):
# Every parent of a node must be an internal node (and not a leaf),
# except for node 0 (the root), whose parent is None.
if pos != 0:
assert isinstance(tree_repair.parent(pos).node, Node)
# Leaf nodes can't have children
if isinstance(node, Leaf):
assert all(c.node is None for c in tree_repair.children(pos))
def test_save_sparseness(n_children):
factory = GraphFactory(31, 1e5, 4)
tree = SBT(factory, d=n_children)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*old_result, sep='\n')
with utils.TempDirectory() as location:
tree.save(os.path.join(location, 'demo'), sparseness=1.0)
tree_loaded = SBT.load(os.path.join(location, 'demo'),
leaf_loader=SigLeaf.load)
assert all(not isinstance(n, Node) for n in tree_loaded.nodes.values())
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {
str(s)
for s in tree_loaded.find(search_minhashes, to_search.data, 0.1)
}
print(*new_result, sep='\n')
assert old_result == new_result
for pos, node in list(tree_loaded.nodes.items()):
# Every parent of a node must be an internal node (and not a leaf),
# except for node 0 (the root), whose parent is None.
if pos != 0:
assert isinstance(tree_loaded.parent(pos).node, Node)
# Leaf nodes can't have children
if isinstance(node, Leaf):
assert all(c.node is None for c in tree_loaded.children(pos))
def test_load_zip_uncompressed(tmpdir):
# uncompress zipped SBT into a tmpdir and search unpacked SBT
import zipfile
testdata = utils.get_test_data("v6.sbt.zip")
testsbt = tmpdir.join("v6.sbt.json")
with zipfile.ZipFile(testdata, 'r') as z:
z.extractall(str(tmpdir))
tree = SBT.load(str(testsbt), leaf_loader=SigLeaf.load)
to_search = load_one_signature(utils.get_test_data(utils.SIG_FILES[0]))
print("*" * 60)
print("{}:".format(to_search))
new_result = {str(s) for s in tree.find(search_minhashes, to_search, 0.1)}
print(*new_result, sep="\n")
assert len(new_result) == 2
def test_sbt_ipfsstorage():
ipfshttpclient = pytest.importorskip('ipfshttpclient')
factory = GraphFactory(31, 1e5, 4)
with utils.TempDirectory() as location:
tree = SBT(factory)
for f in utils.SIG_FILES:
sig = load_one_signature(utils.get_test_data(f))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*old_result, sep='\n')
try:
with IPFSStorage() as storage:
tree.save(os.path.join(location, 'tree'), storage=storage)
except ipfshttpclient.exceptions.ConnectionError:
pytest.xfail("ipfs not installed/functioning probably")
with IPFSStorage() as storage:
tree = SBT.load(os.path.join(location, 'tree'),
leaf_loader=SigLeaf.load,
storage=storage)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*new_result, sep='\n')
assert old_result == new_result
def test_sbt_redisstorage():
redis = pytest.importorskip('redis')
factory = GraphFactory(31, 1e5, 4)
with utils.TempDirectory() as location:
tree = SBT(factory)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*old_result, sep='\n')
try:
with RedisStorage() as storage:
tree.save(os.path.join(location, 'tree'), storage=storage)
except redis.exceptions.ConnectionError:
pytest.xfail("Couldn't connect to redis server")
with RedisStorage() as storage:
tree = SBT.load(os.path.join(location, 'tree'),
leaf_loader=SigLeaf.load,
storage=storage)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*new_result, sep='\n')
assert old_result == new_result
def test_save_sparseness(n_children):
factory = GraphFactory(31, 1e5, 4)
tree = SBT(factory, d=n_children)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {str(s) for s in tree.find(search_minhashes,
to_search.data, 0.1)}
print(*old_result, sep='\n')
with utils.TempDirectory() as location:
tree.save(os.path.join(location, 'demo'), sparseness=1.0)
tree_loaded = SBT.load(os.path.join(location, 'demo'),
leaf_loader=SigLeaf.load)
assert all(not isinstance(n, Node) for n in tree_loaded.nodes.values())
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {str(s) for s in tree_loaded.find(search_minhashes,
to_search.data, 0.1)}
print(*new_result, sep='\n')
assert old_result == new_result
for pos, node in list(tree_loaded.nodes.items()):
# Every parent of a node must be an internal node (and not a leaf),
# except for node 0 (the root), whose parent is None.
if pos != 0:
assert isinstance(tree_loaded.parent(pos).node, Node)
# Leaf nodes can't have children
if isinstance(node, Leaf):
assert all(c.node is None for c in tree_loaded.children(pos))
def test_sbt_zipstorage(tmpdir):
# create tree, save to a zip, then load and search.
factory = GraphFactory(31, 1e5, 4)
tree = SBT(factory)
for f in utils.SIG_FILES:
sig = next(load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*old_result, sep='\n')
with ZipStorage(str(tmpdir.join("tree.sbt.zip"))) as storage:
tree.save(str(tmpdir.join("tree")), storage=storage)
with ZipStorage(str(tmpdir.join("tree.sbt.zip"))) as storage:
tree = SBT.load(str(tmpdir.join("tree")),
leaf_loader=SigLeaf.load,
storage=storage)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*new_result, sep='\n')
assert old_result == new_result
def test_sbt_ipfsstorage():
ipfsapi = pytest.importorskip('ipfsapi')
factory = GraphFactory(31, 1e5, 4)
with utils.TempDirectory() as location:
tree = SBT(factory)
for f in utils.SIG_FILES:
sig = next(signature.load_signatures(utils.get_test_data(f)))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {str(s) for s in tree.find(search_minhashes,
to_search.data, 0.1)}
print(*old_result, sep='\n')
try:
with IPFSStorage() as storage:
tree.save(os.path.join(location, 'tree'), storage=storage)
except ipfsapi.exceptions.ConnectionError:
pytest.xfail("ipfs not installed/functioning probably")
with IPFSStorage() as storage:
tree = SBT.load(os.path.join(location, 'tree'),
leaf_loader=SigLeaf.load,
storage=storage)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {str(s) for s in tree.find(search_minhashes,
to_search.data, 0.1)}
print(*new_result, sep='\n')
assert old_result == new_result
def test_sbt_tarstorage():
factory = GraphFactory(31, 1e5, 4)
with utils.TempDirectory() as location:
tree = SBT(factory)
for f in utils.SIG_FILES:
sig = load_one_signature(utils.get_test_data(f))
leaf = SigLeaf(os.path.basename(f), sig)
tree.add_node(leaf)
to_search = leaf
print('*' * 60)
print("{}:".format(to_search.metadata))
old_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*old_result, sep='\n')
with TarStorage(os.path.join(location, 'tree.tar.gz')) as storage:
tree.save(os.path.join(location, 'tree'), storage=storage)
with TarStorage(os.path.join(location, 'tree.tar.gz')) as storage:
tree = SBT.load(os.path.join(location, 'tree'),
leaf_loader=SigLeaf.load,
storage=storage)
print('*' * 60)
print("{}:".format(to_search.metadata))
new_result = {
str(s)
for s in tree.find(search_minhashes, to_search.data, 0.1)
}
print(*new_result, sep='\n')
assert old_result == new_result
def test_simple(n_children):
factory = GraphFactory(5, 100, 3)
root = SBT(factory, d=n_children)
leaf1 = Leaf("a", factory())
leaf1.data.count('AAAAA')
leaf1.data.count('AAAAT')
leaf1.data.count('AAAAC')
leaf2 = Leaf("b", factory())
leaf2.data.count('AAAAA')
leaf2.data.count('AAAAT')
leaf2.data.count('AAAAG')
leaf3 = Leaf("c", factory())
leaf3.data.count('AAAAA')
leaf3.data.count('AAAAT')
leaf3.data.count('CAAAA')
leaf4 = Leaf("d", factory())
leaf4.data.count('AAAAA')
leaf4.data.count('CAAAA')
leaf4.data.count('GAAAA')
leaf5 = Leaf("e", factory())
leaf5.data.count('AAAAA')
leaf5.data.count('AAAAT')
leaf5.data.count('GAAAA')
root.add_node(leaf1)
root.add_node(leaf2)
root.add_node(leaf3)
root.add_node(leaf4)
root.add_node(leaf5)
def search_kmer(obj, seq):
return obj.data.get(seq)
leaves = [leaf1, leaf2, leaf3, leaf4, leaf5]
kmers = ["AAAAA", "AAAAT", "AAAAG", "CAAAA", "GAAAA"]
def search_kmer_in_list(kmer):
x = []
for l in leaves:
if l.data.get(kmer):
x.append(l)
return set(x)
for kmer in kmers:
assert set(root.find(search_kmer, kmer)) == search_kmer_in_list(kmer)
print('-----')
print([x.metadata for x in root.find(search_kmer, "AAAAA")])
print([x.metadata for x in root.find(search_kmer, "AAAAT")])
print([x.metadata for x in root.find(search_kmer, "AAAAG")])
print([x.metadata for x in root.find(search_kmer, "CAAAA")])
print([x.metadata for x in root.find(search_kmer, "GAAAA")])
with utils.TempDirectory() as location:
root.save(os.path.join(location, 'demo'))
root = SBT.load(os.path.join(location, 'demo'))
for kmer in kmers:
new_result = {str(r) for r in root.find(search_kmer, kmer)}
print(*new_result, sep='\n')
assert new_result == {str(r) for r in search_kmer_in_list(kmer)}
