49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
from spacy.tokens.doc import Doc
|
|
from spacy.tokens.graph import Graph
|
|
from spacy.vocab import Vocab
|
|
|
|
|
|
def test_graph_init():
|
|
doc = Doc(Vocab(), words=["a", "b", "c", "d"])
|
|
graph = Graph(doc, name="hello")
|
|
assert graph.name == "hello"
|
|
assert graph.doc is doc
|
|
|
|
|
|
def test_graph_edges_and_nodes():
|
|
doc = Doc(Vocab(), words=["a", "b", "c", "d"])
|
|
graph = Graph(doc, name="hello")
|
|
node1 = graph.add_node((0,))
|
|
assert graph.get_node((0,)) == node1
|
|
node2 = graph.add_node((1, 3))
|
|
assert list(node2) == [1, 3]
|
|
graph.add_edge(node1, node2, label="one", weight=-10.5)
|
|
assert graph.has_edge(node1, node2, label="one")
|
|
assert node1.heads() == []
|
|
assert [tuple(h) for h in node2.heads()] == [(0,)]
|
|
assert [tuple(t) for t in node1.tails()] == [(1, 3)]
|
|
assert [tuple(t) for t in node2.tails()] == []
|
|
|
|
|
|
def test_graph_walk():
|
|
doc = Doc(Vocab(), words=["a", "b", "c", "d"])
|
|
graph = Graph(
|
|
doc,
|
|
name="hello",
|
|
nodes=[(0,), (1,), (2,), (3,)],
|
|
edges=[(0, 1), (0, 2), (0, 3), (3, 0)],
|
|
labels=None,
|
|
weights=None,
|
|
)
|
|
node0, node1, node2, node3 = list(graph.nodes)
|
|
assert [tuple(h) for h in node0.heads()] == [(3,)]
|
|
assert [tuple(h) for h in node1.heads()] == [(0,)]
|
|
assert [tuple(h) for h in node0.walk_heads()] == [(3,), (0,)]
|
|
assert [tuple(h) for h in node1.walk_heads()] == [(0,), (3,), (0,)]
|
|
assert [tuple(h) for h in node2.walk_heads()] == [(0,), (3,), (0,)]
|
|
assert [tuple(h) for h in node3.walk_heads()] == [(0,), (3,)]
|
|
assert [tuple(t) for t in node0.walk_tails()] == [(1,), (2,), (3,), (0,)]
|
|
assert [tuple(t) for t in node1.walk_tails()] == []
|
|
assert [tuple(t) for t in node2.walk_tails()] == []
|
|
assert [tuple(t) for t in node3.walk_tails()] == [(0,), (1,), (2,), (3,)]
|