Add optional source/target name for Relation

gyorilab · Jul 20, 2024 · d44347c · d44347c
1 parent e864377
commit d44347c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git 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
@@ -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,13 +183,19 @@ 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
         self.target_ns = target_ns
         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