fix bug of accessing non-existent objects in link graph

NPLinker · Jul 15, 2024 · a18ed84 · a18ed84
1 parent a1e9beb
commit a18ed84
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/nplinker/scoring/link_graph.py b/src/nplinker/scoring/link_graph.py
@@ -258,9 +258,13 @@ def filter(self, u_nodes: Sequence[Entity], v_nodes: Sequence[Entity] = [], /) -
     @validate_u
     def _filter_one_node(self, u: Entity, lg: LinkGraph) -> None:
         """Filter the links for a given object and add them to the new LinkGraph object."""
-        links = self[u]
-        for node2, value in links.items():
-            lg.add_link(u, node2, **value)
+        try:
+            links = self[u]
+        except KeyError:
+            pass
+        else:
+            for node2, value in links.items():
+                lg.add_link(u, node2, **value)
 
     @validate_uv
     def _filter_two_nodes(self, u: Entity, v: Entity, lg: LinkGraph) -> None:

diff --git a/tests/unit/scoring/test_link_graph.py b/tests/unit/scoring/test_link_graph.py
@@ -90,8 +90,8 @@ def test_filter(gcfs, spectra, score):
     lg.add_link(gcfs[0], spectra[0], metcalf=score)
     lg.add_link(gcfs[1], spectra[1], metcalf=score)
 
-    u_nodes = [gcfs[0], gcfs[1]]
-    v_nodes = [spectra[0], spectra[1]]
+    u_nodes = [gcfs[0], gcfs[1], gcfs[2]]
+    v_nodes = [spectra[0], spectra[1], spectra[2]]
 
     # test filtering with GCFs
     lg_filtered = lg.filter(u_nodes)