diff --git a/module/base/filter.py b/module/base/filter.py index 2f2332510..f4d3b731f 100644 --- a/module/base/filter.py +++ b/module/base/filter.py @@ -41,34 +41,48 @@ class Filter: def is_preset(self, filter): return len(filter) and filter.lower() in self.preset - def apply(self, objs, func=None): + def apply(self, objs, sort_fn=None, filter_fn=None): """ Args: - objs (list): List of objects and strings - func (callable): A function to filter object. + objs (list): List of objects + filter_fn (callable): A function to filter object. Function should receive an object as arguments, and return a bool. True means add it to output. + sort_fn (callable): A function to sort object. + Function should receive an object as arguments, and return a float/int, the larger the better. + + Returns: list: A list of objects and preset strings, such as [object, object, object, 'reset'] """ out = [] for raw, filter in zip(self.filter_raw, self.filter): - if self.is_preset(raw): + if raw.lower() == 'random': + candidates = objs + if sort_fn is not None: + candidates = list(sorted(candidates, key=sort_fn, reverse=True)) + for c in candidates: + if c not in out: + out.append(c) + elif self.is_preset(raw): raw = raw.lower() if raw not in out: out.append(raw) else: - for index, obj in enumerate(objs): - if self.apply_filter_to_obj(obj=obj, filter=filter) and obj not in out: - out.append(obj) + candidates = [o for o in objs if self.apply_filter_to_obj(o, filter)] + if sort_fn is not None: + candidates = list(sorted(candidates, key=sort_fn, reverse=True)) + for c in candidates: + if c not in out: + out.append(c) - if func is not None: + if filter_fn is not None: objs, out = out, [] for obj in objs: if isinstance(obj, str): out.append(obj) - elif func(obj): + elif filter_fn(obj): out.append(obj) else: # Drop this object diff --git a/tasks/rogue/blessing/blessing.py b/tasks/rogue/blessing/blessing.py index a83db202f..550223f41 100644 --- a/tasks/rogue/blessing/blessing.py +++ b/tasks/rogue/blessing/blessing.py @@ -12,6 +12,7 @@ from tasks.rogue.assets.assets_rogue_blessing import * from tasks.rogue.assets.assets_rogue_ui import BLESSING_CONFIRM from tasks.rogue.blessing.preset import BLESSING_PRESET, RESONANCE_PRESET from tasks.rogue.blessing.selector import RogueSelector +from tasks.rogue.blessing.ui import RogueUI from tasks.rogue.blessing.utils import get_regex_from_keyword_name, is_card_selected, parse_name from tasks.rogue.keywords import * @@ -108,6 +109,10 @@ class RogueBlessingSelector(RogueSelector): self.recognize_and_select() """ + def __init__(self, main: RogueUI): + super().__init__(main) + self.sort_fn = lambda _: _.rarity + def get_blessing_count(self) -> int: """ Returns: The number of blessing diff --git a/tasks/rogue/blessing/selector.py b/tasks/rogue/blessing/selector.py index fa7bcc0da..68784cc67 100644 --- a/tasks/rogue/blessing/selector.py +++ b/tasks/rogue/blessing/selector.py @@ -15,6 +15,7 @@ class RogueSelector: self.main = main self.filter_ = None self.ocr_results = [] + self.sort_fn = None def recognition(self): ... @@ -55,8 +56,8 @@ class RogueSelector: return None if self.filter_: - keywords = [result.matched_keyword for result in self.ocr_results] - priority = self.filter_.apply(keywords) + keywords = [result.matched_keyword for result in self.ocr_results if result.matched_keyword is not None] + priority = self.filter_.apply(keywords, self.sort_fn) priority = [option if isinstance(option, str) else match_ocr_result(option) for option in priority] else: logger.warning("No filter loaded, use random instead")