diff --git a/nebula3/data/DataObject.py b/nebula3/data/DataObject.py index 838ac270..fa60c5d9 100644 --- a/nebula3/data/DataObject.py +++ b/nebula3/data/DataObject.py @@ -710,7 +710,10 @@ def _cast_node(node: Vertex): return { "vid": node.get_id().cast(), "tags": { - tag_name: node.properties(tag_name) for tag_name in node.tags() + tag_name: { + k: v.cast() for k, v in node.properties(tag_name).items() + } + for tag_name in node.tags() }, } @@ -720,7 +723,7 @@ def _cast_relationship(edge: Edge): "dst": edge.end_vertex_id().cast(), "type": edge.edge_name(), "rank": edge.ranking(), - "props": edge.properties(), + "props": {k: v.cast() for k, v in edge.properties().items()}, } def _cast_primitive(raw_value): diff --git a/tests/test_data_type.py b/tests/test_data_type.py index 1e3dc321..41f41e4e 100644 --- a/tests/test_data_type.py +++ b/tests/test_data_type.py @@ -386,22 +386,6 @@ def test_cast(self): def test_cast_primitive(self): # Test casting for primitive types - def _cast_node(node: Vertex): - return { - "vid": node.get_id().cast(), - "tags": { - tag_name: node.properties(tag_name) for tag_name in node.tags() - }, - } - - def _cast_relationship(edge: Edge): - return { - "src": edge.start_vertex_id().cast(), - "dst": edge.end_vertex_id().cast(), - "type": edge.edge_name(), - "rank": edge.ranking(), - "props": edge.properties(), - } # Test boolean bool_val = ttypes.Value(bVal=True) @@ -423,37 +407,6 @@ def _cast_relationship(edge: Edge): null_val = ttypes.Value(nVal=ttypes.NullType.__NULL__) assert ValueWrapper(null_val).cast_primitive() is None - # Test node - node_val = ttypes.Value(vVal=self.get_vertex_value(b"Tom")) - node = ValueWrapper(node_val).as_node() - assert ValueWrapper(node_val).cast_primitive() == { - "vid": node.get_id().cast(), - "tags": {tag_name: node.properties(tag_name) for tag_name in node.tags()}, - } - - # Test relationship - relationship_val = ttypes.Value(eVal=self.get_edge_value(b"Tom", b"Lily")) - edge = ValueWrapper(relationship_val).as_relationship() - assert ValueWrapper(relationship_val).cast_primitive() == { - "src": edge.start_vertex_id().cast(), - "dst": edge.end_vertex_id().cast(), - "type": edge.edge_name(), - "rank": edge.ranking(), - "props": edge.properties(), - } - - # Test path - path_val = ttypes.Value(pVal=self.get_path_value(b"Tom")) - path_raw = ValueWrapper(path_val) - path = path_raw.as_path() - path_primitive = path_raw.cast_primitive() - assert path_primitive == { - "path_str": path.__repr__(), - "start_node": _cast_node(path.start_node()), - "edges": [_cast_relationship(x) for x in path.relationships()], - "nodes": [_cast_node(x) for x in path.nodes()], - } - # Test geography geography_val = ttypes.Value(ggVal=self.get_geography_value(3.0, 5.2)) geography_raw = ValueWrapper(geography_val)