Skip to content

Commit c9ce0b5

Browse files
adamnschFlorentinD
adamnsch
and
FlorentinD
committed
WIP: Move custom fields into properties dict
Co-Authored-By: Florentin Dörre <florentin.dorre@neotechnology.com>
1 parent d0dae64 commit c9ce0b5

File tree

6 files changed

+68
-74
lines changed

6 files changed

+68
-74
lines changed

‎python-wrapper/src/neo4j_viz/neo4j.py‎

Lines changed: 8 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -79,13 +79,10 @@ def _map_node(node: neo4j.graph.Node, size_property: Optional[str], caption_prop
7979
else:
8080
caption = str(node.get(caption_property))
8181

82-
base_node_props = dict(id=node.element_id, caption=caption, labels=labels, size=size)
82+
properties = {k: v for k, v in node.items()}
83+
properties["__labels"] = labels
8384

84-
protected_props = base_node_props.keys()
85-
additional_node_props = {k: v for k, v in node.items()}
86-
additional_node_props = _rename_protected_props(additional_node_props, protected_props)
87-
88-
return Node(**base_node_props, **additional_node_props)
85+
return Node(id=node.element_id, caption=caption, size=size, properties=properties)
8986

9087

9188
def _map_relationship(rel: neo4j.graph.Relationship, caption_property: Optional[str]) -> Optional[Relationship]:
@@ -100,32 +97,13 @@ def _map_relationship(rel: neo4j.graph.Relationship, caption_property: Optional[
10097
else:
10198
caption = None
10299

103-
base_rel_props = dict(
100+
properties = {k: v for k, v in rel.items()}
101+
properties["__type"] = rel.type
102+
103+
return Relationship(
104104
id=rel.element_id,
105105
source=rel.start_node.element_id,
106106
target=rel.end_node.element_id,
107-
_type=rel.type,
108107
caption=caption,
108+
properties=properties,
109109
)
110-
111-
protected_props = base_rel_props.keys()
112-
additional_rel_props = {k: v for k, v in rel.items()}
113-
additional_rel_props = _rename_protected_props(additional_rel_props, protected_props)
114-
115-
return Relationship(
116-
**base_rel_props,
117-
**additional_rel_props,
118-
)
119-
120-
121-
def _rename_protected_props(
122-
additional_props: dict[str, Any],
123-
protected_props: Iterable[str],
124-
) -> dict[str, Union[str, int, float]]:
125-
for prop in protected_props:
126-
if prop not in additional_props:
127-
continue
128-
129-
additional_props[f"__{prop}"] = additional_props.pop(prop)
130-
131-
return additional_props

‎python-wrapper/src/neo4j_viz/node.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ class Node(BaseModel, extra="allow"):
4545
x: Optional[RealNumber] = Field(None, description="The x-coordinate of the node")
4646
#: The y-coordinate of the node
4747
y: Optional[RealNumber] = Field(None, description="The y-coordinate of the node")
48+
#: The properties of the node
49+
properties: dict[str, Any] = Field(default_factory=dict, description="The properties of the node")
4850

4951
@field_serializer("color")
5052
def serialize_color(self, color: Color) -> str:

‎python-wrapper/src/neo4j_viz/relationship.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ class Relationship(BaseModel, extra="allow"):
4343
)
4444
#: The color of the relationship. Allowed input is for example "#FF0000", "red" or (255, 0, 0)
4545
color: Optional[ColorType] = Field(None, description="The color of the relationship")
46+
#: The properties of the relationship
47+
properties: dict[str, Any] = Field(default_factory=dict, description="The properties of the relationship")
4648

4749
@field_serializer("color")
4850
def serialize_color(self, color: Color) -> str:

‎python-wrapper/tests/test_neo4j.py‎

Lines changed: 45 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -31,27 +31,29 @@ def test_from_neo4j_graph(neo4j_session: Session) -> None:
3131
Node(
3232
id=node_ids[0],
3333
caption="_CI_A",
34-
labels=["_CI_A"],
35-
name="Alice",
36-
height=20,
37-
__id=42,
38-
_id=1337,
39-
__caption="hello",
34+
properties=dict(
35+
__labels=["_CI_A"],
36+
name="Alice",
37+
height=20,
38+
_id=1337,
39+
caption="hello",
40+
)
4041
),
4142
Node(
4243
id=node_ids[1],
4344
caption="_CI_A:_CI_B",
44-
labels=["_CI_A", "_CI_B"],
45-
name="Bob",
46-
height=10,
47-
__id=84,
48-
__size=11,
49-
__labels=[1, 2],
45+
properties=dict(
46+
__labels=["_CI_A", "_CI_B"],
47+
name="Bob",
48+
height=10,
49+
size=11,
50+
labels=[1, 2],
51+
),
5052
),
5153
]
5254

5355
assert len(VG.nodes) == 2
54-
assert sorted(VG.nodes, key=lambda x: x.name) == expected_nodes # type: ignore[attr-defined]
56+
assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes # type: ignore[attr-defined]
5557

5658
assert len(VG.relationships) == 2
5759
vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")
@@ -76,27 +78,29 @@ def test_from_neo4j_result(neo4j_session: Session) -> None:
7678
Node(
7779
id=node_ids[0],
7880
caption="_CI_A",
79-
labels=["_CI_A"],
80-
name="Alice",
81-
height=20,
82-
__id=42,
83-
_id=1337,
84-
__caption="hello",
81+
properties=dict(
82+
__labels=["_CI_A"],
83+
name="Alice",
84+
height=20,
85+
_id=1337,
86+
caption="hello",
87+
)
8588
),
8689
Node(
8790
id=node_ids[1],
8891
caption="_CI_A:_CI_B",
89-
labels=["_CI_A", "_CI_B"],
90-
name="Bob",
91-
height=10,
92-
__id=84,
93-
__size=11,
94-
__labels=[1, 2],
92+
properties=dict(
93+
__labels=["_CI_A", "_CI_B"],
94+
name="Bob",
95+
height=10,
96+
size=11,
97+
labels=[1, 2],
98+
)
9599
),
96100
]
97101

98102
assert len(VG.nodes) == 2
99-
assert sorted(VG.nodes, key=lambda x: x.name) == expected_nodes # type: ignore[attr-defined]
103+
assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes # type: ignore[attr-defined]
100104

101105
assert len(VG.relationships) == 2
102106
vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")
@@ -119,29 +123,31 @@ def test_from_neo4j_graph_full(neo4j_session: Session) -> None:
119123
Node(
120124
id=node_ids[0],
121125
caption="Alice",
122-
labels=["_CI_A"],
123-
name="Alice",
124-
height=20,
125126
size=60.0,
126-
__id=42,
127-
_id=1337,
128-
__caption="hello",
127+
properties=dict(
128+
__labels=["_CI_A"],
129+
name="Alice",
130+
height=20,
131+
_id=1337,
132+
caption="hello",
133+
)
129134
),
130135
Node(
131136
id=node_ids[1],
132137
caption="Bob",
133-
labels=["_CI_A", "_CI_B"],
134-
name="Bob",
135-
height=10,
136138
size=3.0,
137-
__id=84,
138-
__size=11,
139-
__labels=[1, 2],
139+
properties=dict(
140+
__labels=["_CI_A", "_CI_B"],
141+
name="Bob",
142+
height=10,
143+
size=11,
144+
labels=[1, 2],
145+
),
140146
),
141147
]
142148

143149
assert len(VG.nodes) == 2
144-
assert sorted(VG.nodes, key=lambda x: x.name) == expected_nodes # type: ignore[attr-defined]
150+
assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes # type: ignore[attr-defined]
145151

146152
assert len(VG.relationships) == 2
147153
vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")

‎python-wrapper/tests/test_node.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ def test_nodes_with_all_options() -> None:
2626
"pinned": True,
2727
"x": 1,
2828
"y": 10,
29+
"properties": {},
2930
}
3031

3132

@@ -36,6 +37,7 @@ def test_nodes_minimal_node() -> None:
3637

3738
assert node.to_dict() == {
3839
"id": "1",
40+
"properties": {},
3941
}
4042

4143

@@ -48,6 +50,7 @@ def test_node_with_float_size() -> None:
4850
assert node.to_dict() == {
4951
"id": "1",
5052
"size": 10.2,
53+
"properties": {},
5154
}
5255

5356

@@ -60,6 +63,7 @@ def test_node_with_additional_fields() -> None:
6063
assert node.to_dict() == {
6164
"id": "1",
6265
"componentId": 2,
66+
"properties": {},
6367
}
6468

6569

@@ -69,6 +73,7 @@ def test_id_aliases(alias: str) -> None:
6973

7074
assert node.to_dict() == {
7175
"id": "1",
76+
"properties": {},
7277
}
7378

7479

‎python-wrapper/tests/test_relationship.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ def test_rels_with_all_options() -> None:
2323
"captionAlign": "top",
2424
"captionSize": 12,
2525
"color": "#ff0000",
26+
"properties": {},
2627
}
2728

2829

@@ -34,7 +35,7 @@ def test_rels_minimal_rel() -> None:
3435

3536
rel_dict = rel.to_dict()
3637

37-
assert {"id", "from", "to"} == set(rel_dict.keys())
38+
assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
3839
assert rel_dict["from"] == "1"
3940
assert rel_dict["to"] == "2"
4041

@@ -43,12 +44,12 @@ def test_rels_additional_fields() -> None:
4344
rel = Relationship(
4445
source="1",
4546
target="2",
46-
componentId=2,
47+
properties=dict(componentId=2),
4748
)
4849

4950
rel_dict = rel.to_dict()
50-
assert {"id", "from", "to", "componentId"} == set(rel_dict.keys())
51-
assert rel.componentId == 2 # type: ignore[attr-defined]
51+
assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
52+
assert rel.properties["componentId"] == 2 # type: ignore[attr-defined]
5253

5354

5455
@pytest.mark.parametrize("src_alias", ["source", "sourceNodeId", "source_node_id", "from"])
@@ -63,6 +64,6 @@ def test_aliases(src_alias: str, trg_alias: str) -> None:
6364

6465
rel_dict = rel.to_dict()
6566

66-
assert {"id", "from", "to"} == set(rel_dict.keys())
67+
assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
6768
assert rel_dict["from"] == "1"
6869
assert rel_dict["to"] == "2"

0 commit comments

Comments
 (0)