From 04853b6c3190c202a2f84083a1523f1a4e73ca97 Mon Sep 17 00:00:00 2001
From: YoursFunny <admin@yoursfunny.top>
Date: Tue, 21 Nov 2023 19:53:02 +0800
Subject: [PATCH] feat: support load search for buttons

---
 module/base/button.py | 32 ++++++++++++++++++++++++++------
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/module/base/button.py b/module/base/button.py
index d0bf122..e3ac323 100644
--- a/module/base/button.py
+++ b/module/base/button.py
@@ -144,16 +144,25 @@ class ButtonWrapper(Resource):
     def __bool__(self):
         return True
 
+    def iter_buttons(self) -> t.Iterator[Button]:
+        for _, assets in self.data_buttons.items():
+            if isinstance(assets, Button):
+                yield assets
+            elif isinstance(assets, list):
+                for asset in assets:
+                    yield asset
+
     @cached_property
     def buttons(self) -> t.List[Button]:
-        # for trial in [server.lang, 'share', 'cn']:
-        for trial in [server.lang, 'share', 'jp']:
-            assets = self.data_buttons.get(trial, None)
-            if assets is not None:
+        for trial in [server.lang, 'share', 'cn']:
+            try:
+                assets = self.data_buttons[trial]
                 if isinstance(assets, Button):
                     return [assets]
                 elif isinstance(assets, list):
                     return assets
+            except KeyError:
+                pass
 
         raise ScriptError(f'ButtonWrapper({self}) on server {server.lang} has no fallback button')
 
@@ -222,13 +231,24 @@ class ButtonWrapper(Resource):
         """
         if isinstance(button, ButtonWrapper):
             button = button.matched_button
-        for b in self.buttons:
+        for b in self.iter_buttons():
             b.load_offset(button)
 
     def clear_offset(self):
-        for b in self.buttons:
+        for b in self.iter_buttons():
             b.clear_offset()
 
+    def load_search(self, area):
+        """
+        Set `search` attribute.
+        Note that this method is irreversible.
+
+        Args:
+            area:
+        """
+        for b in self.iter_buttons():
+            b.search = area
+
 
 class ClickButton:
     def __init__(self, area, button=None, name='CLICK_BUTTON'):