|
1
|
|
|
import networkx as nx |
|
2
|
|
|
|
|
3
|
|
|
from networkx_query import prepare_query, search_edges, search_nodes |
|
4
|
|
|
|
|
5
|
|
|
g = nx.DiGraph() |
|
6
|
|
|
g.add_node(1, product="chocolate") |
|
7
|
|
|
g.add_node(2, product="milk") |
|
8
|
|
|
g.add_node(3, product="coat") |
|
9
|
|
|
g.add_edge(1, 2, action="shake") |
|
10
|
|
|
g.add_edge(3, 2, action="produce") |
|
11
|
|
|
|
|
12
|
|
|
|
|
13
|
|
|
def test_prepare_query(): |
|
14
|
|
|
predicate = prepare_query({"eq": [("product",), "chocolate"]}) |
|
15
|
|
|
assert predicate({'product': 'chocolate'}) |
|
16
|
|
|
assert not predicate({'product': 'milk'}) |
|
17
|
|
|
|
|
18
|
|
|
|
|
19
|
|
|
def test_search_nodes_implementation(): |
|
20
|
|
|
predicate = prepare_query({"eq": [("product",), "chocolate"]}) |
|
21
|
|
|
assert [(node[0], predicate(node[1])) for node in g.nodes(data=True)] == [(1, True), (2, False), (3, False)] |
|
22
|
|
|
|
|
23
|
|
|
|
|
24
|
|
|
def test_search_edges_implementation(): |
|
25
|
|
|
predicate = prepare_query({"eq": [("action",), "shake"]}) |
|
26
|
|
|
assert [((edge[0], edge[1]), predicate(edge[2])) for edge in g.edges(data=True)] == [ |
|
27
|
|
|
((1, 2), True), |
|
28
|
|
|
((3, 2), False), |
|
29
|
|
|
] |
|
30
|
|
|
|
|
31
|
|
|
|
|
32
|
|
|
def test_search_nodes(): |
|
33
|
|
|
assert list(search_nodes(g, {"eq": [("product",), "chocolate"]})) == [1] |
|
34
|
|
|
assert list(search_nodes(g, {"eq": [("product",), "milk"]})) == [2] |
|
35
|
|
|
|
|
36
|
|
|
result = list(search_nodes(g, {"in": [("product",), ["milk", "coat"]]})) |
|
37
|
|
|
assert len(result) == 2 |
|
38
|
|
|
assert 2 in result |
|
39
|
|
|
assert 3 in result |
|
40
|
|
|
|
|
41
|
|
|
assert list(search_nodes(g, {"eq": [("product",), "none"]})) == [] |
|
42
|
|
|
|
|
43
|
|
|
|
|
44
|
|
|
def test_search_edges(): |
|
45
|
|
|
assert list(search_edges(g, {"eq": [("action",), "produce"]})) == [(3, 2)] |
|
46
|
|
|
assert list(search_edges(g, {"eq": [("action",), "shake"]})) == [(1, 2)] |
|
47
|
|
|
assert list(search_edges(g, {"eq": [("action",), "none"]})) == [] |
|
48
|
|
|
|