diff --git a/grandcypher/test_types.py b/grandcypher/test_types.py new file mode 100644 index 0000000..710dc3b --- /dev/null +++ b/grandcypher/test_types.py @@ -0,0 +1,54 @@ +import pickle + +import networkx as nx + +from grandcypher import GrandCypher +from grandcypher.types import AttributeRef, IDRef, EntityRef + + +def test_attribute_ref_pickle_roundtrip(): + ref = AttributeRef('a', 'name') + restored = pickle.loads(pickle.dumps(ref)) + assert restored == 'a.name' + assert type(restored) is AttributeRef + assert restored.entity_name == 'a' + assert restored.attribute == 'name' + + +def test_id_ref_pickle_roundtrip(): + ref = IDRef('a') + restored = pickle.loads(pickle.dumps(ref)) + assert restored == 'ID(a)' + assert type(restored) is IDRef + assert restored.entity_name == 'a' + + +def test_entity_ref_pickle_roundtrip(): + ref = EntityRef('a') + restored = pickle.loads(pickle.dumps(ref)) + assert restored == 'a' + assert type(restored) is EntityRef + assert restored.entity_name == 'a' + + +def test_query_results_pickle_roundtrip(): + host = nx.DiGraph() + host.add_node("x", name="Alice", age=30) + host.add_node("y", name="Bob", age=25) + host.add_edge("x", "y", since=2020) + + qry = """ + MATCH (A)-[E]->(B) + WHERE A.name == "Alice" + RETURN A, A.name, A.age, ID(A), B + """ + + results = GrandCypher(host).run(qry) + restored = pickle.loads(pickle.dumps(results)) + + assert restored == results + for key in results: + assert type(restored[key]) is type(results[key]) + for key in restored: + restored_key_type = type(key) + assert any(type(k) is restored_key_type for k in results) diff --git a/grandcypher/types.py b/grandcypher/types.py index 5754249..777be6a 100644 --- a/grandcypher/types.py +++ b/grandcypher/types.py @@ -12,6 +12,9 @@ def __new__(cls, entity_name): instance.entity_name = str(entity_name) return instance + def __getnewargs__(self): + return (self.entity_name,) + class AttributeRef(str): """Node/edge attribute reference, e.g. A.age. @@ -25,6 +28,9 @@ def __new__(cls, entity_name, attribute): instance.attribute = str(attribute) return instance + def __getnewargs__(self): + return (self.entity_name, self.attribute) + class IDRef(str): """Reference to ID(A) in WHERE clauses. @@ -36,3 +42,6 @@ def __new__(cls, entity_name): instance = super().__new__(cls, f"ID({entity_name})") instance.entity_name = str(entity_name) return instance + + def __getnewargs__(self): + return (self.entity_name,)