| 34 | |
| 35 | |
| 36 | class TranslationCache: |
| 37 | @staticmethod |
| 38 | def _sort_dict_recursively(obj): |
| 39 | if isinstance(obj, dict): |
| 40 | return { |
| 41 | k: TranslationCache._sort_dict_recursively(v) |
| 42 | for k in sorted(obj.keys()) |
| 43 | for v in [obj[k]] |
| 44 | } |
| 45 | elif isinstance(obj, list): |
| 46 | return [TranslationCache._sort_dict_recursively(item) for item in obj] |
| 47 | return obj |
| 48 | |
| 49 | def __init__(self, translate_engine: str, translate_engine_params: dict = None): |
| 50 | assert ( |
| 51 | len(translate_engine) < 20 |
| 52 | ), "current cache require translate engine name less than 20 characters" |
| 53 | self.translate_engine = translate_engine |
| 54 | self.replace_params(translate_engine_params) |
| 55 | |
| 56 | # The program typically starts multi-threaded translation |
| 57 | # only after cache parameters are fully configured, |
| 58 | # so thread safety doesn't need to be considered here. |
| 59 | def replace_params(self, params: dict = None): |
| 60 | if params is None: |
| 61 | params = {} |
| 62 | self.params = params |
| 63 | params = self._sort_dict_recursively(params) |
| 64 | self.translate_engine_params = json.dumps(params) |
| 65 | |
| 66 | def update_params(self, params: dict = None): |
| 67 | if params is None: |
| 68 | params = {} |
| 69 | self.params.update(params) |
| 70 | self.replace_params(self.params) |
| 71 | |
| 72 | def add_params(self, k: str, v): |
| 73 | self.params[k] = v |
| 74 | self.replace_params(self.params) |
| 75 | |
| 76 | # Since peewee and the underlying sqlite are thread-safe, |
| 77 | # get and set operations don't need locks. |
| 78 | def get(self, original_text: str) -> Optional[str]: |
| 79 | result = _TranslationCache.get_or_none( |
| 80 | translate_engine=self.translate_engine, |
| 81 | translate_engine_params=self.translate_engine_params, |
| 82 | original_text=original_text, |
| 83 | ) |
| 84 | return result.translation if result else None |
| 85 | |
| 86 | def set(self, original_text: str, translation: str): |
| 87 | try: |
| 88 | _TranslationCache.create( |
| 89 | translate_engine=self.translate_engine, |
| 90 | translate_engine_params=self.translate_engine_params, |
| 91 | original_text=original_text, |
| 92 | translation=translation, |
| 93 | ) |