python-graph-visualization/python-wrapper/src/neo4j_viz/neo4j.py at ba3382430a0bb6c5bca70ca222b037a45cec9700 · neo4j/python-graph-visualization · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from __future__ import annotations

from collections.abc import Iterable
from typing import Any, Optional, Union

import neo4j.graph
from neo4j import Result

from neo4j_viz.node import Node
from neo4j_viz.relationship import Relationship
from neo4j_viz.visualization_graph import VisualizationGraph


def from_neo4j(
    result: Union[neo4j.graph.Graph, Result],
    size_property: Optional[str] = None,
    node_caption: Optional[str] = "labels",
    relationship_caption: Optional[str] = "type",
    node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
) -> VisualizationGraph:
    """
    Create a VisualizationGraph from a Neo4j Graph or Neo4j Result object.

    All node and relationship properties will be included in the visualization graph.
    If the property names are conflicting with those of `Node` and `Relationship` objects, they will be prefixed
    with `__`.

    Parameters
    ----------
    result : Union[neo4j.graph.Graph, Result]
        Query result either in shape of a Graph or result.
    size_property : str, optional
        Property to use for node size, by default None.
    node_caption : str, optional
        Property to use as the node caption, by default the node labels will be used.
    relationship_caption : str, optional
        Property to use as the relationship caption, by default the relationship type will be used.
    node_radius_min_max : tuple[float, float], optional
        Minimum and maximum node radius, by default (3, 60).
        To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range.
    """

    if isinstance(result, Result):
        graph = result.graph()
    elif isinstance(result, neo4j.graph.Graph):
        graph = result
    else:
        raise ValueError(f"Invalid input type `{type(result)}`. Expected `neo4j.Graph` or `neo4j.Result`")

    nodes = [_map_node(node, size_property, caption_property=node_caption) for node in graph.nodes]
    relationships = []
    for rel in graph.relationships:
        mapped_rel = _map_relationship(rel, caption_property=relationship_caption)
        if mapped_rel:
            relationships.append(mapped_rel)

    VG = VisualizationGraph(nodes, relationships)

    if node_radius_min_max and size_property is not None:
        VG.resize_nodes(node_radius_min_max=node_radius_min_max)

    return VG


def _map_node(node: neo4j.graph.Node, size_property: Optional[str], caption_property: Optional[str]) -> Node:
    labels = sorted([label for label in node.labels])

    if size_property:
        size = node.get(size_property)
    else:
        size = None

    if caption_property:
        if caption_property == "labels":
            if len(labels) > 0:
                caption = ":".join([label for label in labels])
            else:
                caption = None
        else:
            caption = str(node.get(caption_property))

    base_node_props = dict(id=node.element_id, caption=caption, labels=labels, size=size)

    protected_props = base_node_props.keys()
    additional_node_props = {k: v for k, v in node.items()}
    additional_node_props = _rename_protected_props(additional_node_props, protected_props)

    return Node(**base_node_props, **additional_node_props)


def _map_relationship(rel: neo4j.graph.Relationship, caption_property: Optional[str]) -> Optional[Relationship]:
    if rel.start_node is None or rel.end_node is None:
        return None

    if caption_property:
        if caption_property == "type":
            caption = rel.type
        else:
            caption = str(rel.get(caption_property))
    else:
        caption = None

    base_rel_props = dict(
        id=rel.element_id,
        source=rel.start_node.element_id,
        target=rel.end_node.element_id,
        _type=rel.type,
        caption=caption,
    )

    protected_props = base_rel_props.keys()
    additional_rel_props = {k: v for k, v in rel.items()}
    additional_rel_props = _rename_protected_props(additional_rel_props, protected_props)

    return Relationship(
        **base_rel_props,
        **additional_rel_props,
    )


def _rename_protected_props(
    additional_props: dict[str, Any],
    protected_props: Iterable[str],
) -> dict[str, Union[str, int, float]]:
    for prop in protected_props:
        if prop not in additional_props:
            continue

        additional_props[f"__{prop}"] = additional_props.pop(prop)

    return additional_props