From e7854e32eaa2df3788810f0d48e011276ab47dca Mon Sep 17 00:00:00 2001 From: Sebastian Lobentanzer Date: Wed, 29 Nov 2023 10:59:48 +0100 Subject: [PATCH] expand lists when adding entities .. after relationship selection --- biochatter/prompts.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/biochatter/prompts.py b/biochatter/prompts.py index 8c62ec69..b7bd1132 100644 --- a/biochatter/prompts.py +++ b/biochatter/prompts.py @@ -372,16 +372,14 @@ def _select_relationships(self) -> bool: if self.selected_relationship_labels: for key, value in self.selected_relationship_labels.items(): - if ( - value["source"] - and value["source"] not in self.selected_entities - ): - self.selected_entities.append(value["source"]) - if ( - value["target"] - and value["target"] not in self.selected_entities - ): - self.selected_entities.append(value["target"]) + sources = ensure_iterable(value["source"]) + targets = ensure_iterable(value["target"]) + for source in sources: + if source not in self.selected_entities: + self.selected_entities.append(source) + for target in targets: + if target not in self.selected_entities: + self.selected_entities.append(target) return bool(result)