From d44347c6572b5593bda25b7bc2cf27b702416b36 Mon Sep 17 00:00:00 2001
From: Ben Gyori <ben.gyori@gmail.com>
Date: Fri, 19 Jul 2024 21:41:04 -0400
Subject: [PATCH] Add optional source/target name for Relation

---
 src/indra_cogex/client/neo4j_client.py |  5 ++++-
 src/indra_cogex/representation.py      | 10 ++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/indra_cogex/client/neo4j_client.py b/src/indra_cogex/client/neo4j_client.py
index 5a3f22f1d..a03ccd4b2 100644
--- a/src/indra_cogex/client/neo4j_client.py
+++ b/src/indra_cogex/client/neo4j_client.py
@@ -868,8 +868,11 @@ def neo4j_to_relations(neo4j_path: neo4j.graph.Path) -> List[Relation]:
             rel_type = neo4j_relation.type
             props = dict(neo4j_relation)
             source_ns, source_id = process_identifier(neo4j_relation.start_node["id"])
+            source_name = neo4j_relation.start_node.get("name")
             target_ns, target_id = process_identifier(neo4j_relation.end_node["id"])
-            rel = Relation(source_ns, source_id, target_ns, target_id, rel_type, props)
+            target_name = neo4j_relation.end_node.get("name")
+            rel = Relation(source_ns, source_id, target_ns, target_id, rel_type, props,
+                           source_name=source_name, target_name=target_name)
             relations.append(rel)
         return relations
 
diff --git a/src/indra_cogex/representation.py b/src/indra_cogex/representation.py
index 7e3e76786..2cb3a621b 100644
--- a/src/indra_cogex/representation.py
+++ b/src/indra_cogex/representation.py
@@ -164,6 +164,8 @@ def __init__(
         target_id: str,
         rel_type: str,
         data: Optional[Mapping[str, Any]] = None,
+        source_name: Optional[str] = None,
+        target_name: Optional[str] = None,
     ):
         """Initialize the relation.
 
@@ -181,6 +183,10 @@ def __init__(
             The type of relation.
         data :
             An optional data dictionary associated with the relation.
+        source_name :
+            An optional name for the source node.
+        target_name :
+            An optional name for the target node.
         """
         self.source_ns = source_ns
         self.source_id = source_id
@@ -188,6 +194,8 @@ def __init__(
         self.target_id = target_id
         self.rel_type = rel_type
         self.data = data if data else {}
+        self.source_name = source_name
+        self.target_name = target_name
 
     def to_json(self) -> RelJson:
         """Serialize the relation to JSON format.
@@ -204,6 +212,8 @@ def to_json(self) -> RelJson:
             "target_id": self.target_id,
             "rel_type": self.rel_type,
             "data": self.data,
+            "source_name": self.source_name,
+            "target_name": self.target_name,
         }
 
     def __str__(self):  # noqa:D105