spaCy/spacy/tests/doc/test_graph.py
Daniël de Kok e2b70df012
Configure isort to use the Black profile, recursively isort the spacy module (#12721)
* Use isort with Black profile

* isort all the things

* Fix import cycles as a result of import sorting

* Add DOCBIN_ALL_ATTRS type definition

* Add isort to requirements

* Remove isort from build dependencies check

* Typo
2023-06-14 17:48:41 +02:00

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,)]