From 4481da676c20d440abf778b50b1c84883f32a35b Mon Sep 17 00:00:00 2001
From: Zebartin <16185081+Zebartin@users.noreply.github.com>
Date: Wed, 27 Sep 2023 10:53:37 +0800
Subject: [PATCH] Fix: keyword_class in draggable list

---
 module/ui/draggable_list.py | 8 +++-----
 tasks/assignment/ui.py      | 4 ++++
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/module/ui/draggable_list.py b/module/ui/draggable_list.py
index 806ff35b0..7f8f46c92 100644
--- a/module/ui/draggable_list.py
+++ b/module/ui/draggable_list.py
@@ -42,11 +42,9 @@ class DraggableList:
         self.name = name
         self.keyword_class = keyword_class
         self.ocr_class = ocr_class
-        if not isinstance(keyword_class, list):
-            keyword_class = [keyword_class]
-        self.known_rows = [
-            kw for kc in keyword_class for kw in kc.instances.values()
-        ]
+        if isinstance(keyword_class, list):
+            keyword_class = keyword_class[0]
+        self.known_rows = list(keyword_class.instances.values())
         self.search_button = search_button
         self.check_row_order = check_row_order
         self.active_color = active_color
diff --git a/tasks/assignment/ui.py b/tasks/assignment/ui.py
index 8aef740a0..d184c2d5b 100644
--- a/tasks/assignment/ui.py
+++ b/tasks/assignment/ui.py
@@ -121,6 +121,10 @@ ASSIGNMENT_ENTRY_LIST = DraggableList(
     check_row_order=False,
     active_color=(40, 40, 40)
 )
+ASSIGNMENT_ENTRY_LIST.known_rows = [
+    kw for kc in ASSIGNMENT_ENTRY_LIST.keyword_class
+    for kw in kc.instances.values()
+]
 
 
 class AssignmentUI(UI):