autocorrect_data.py (13481B)
1 # Copyright 2021 Google LLC 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # https://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 """Python program to make autocorrect_data.h. 15 This program reads from a prepared dictionary file and generates a C source file 16 "autocorrect_data.h" with a serialized trie embedded as an array. Run this 17 program and pass it as the first argument like: 18 $ qmk generate-autocorrect-data autocorrect_dict.txt 19 Each line of the dict file defines one typo and its correction with the syntax 20 "typo -> correction". Blank lines or lines starting with '#' are ignored. 21 Example: 22 :thier -> their 23 fitler -> filter 24 lenght -> length 25 ouput -> output 26 widht -> width 27 For full documentation, see QMK Docs 28 """ 29 30 import textwrap 31 from typing import Any, Dict, Iterator, List, Tuple 32 33 from milc import cli 34 35 from qmk.commands import dump_lines 36 from qmk.constants import GPL2_HEADER_C_LIKE, GENERATED_HEADER_C_LIKE 37 from qmk.keyboard import keyboard_completer, keyboard_folder 38 from qmk.keymap import keymap_completer, locate_keymap 39 from qmk.path import normpath 40 from qmk.util import maybe_exit 41 42 KC_A = 4 43 KC_SPC = 0x2c 44 KC_QUOT = 0x34 45 46 TYPO_CHARS = dict([ 47 ("'", KC_QUOT), 48 (':', KC_SPC), # "Word break" character. 49 ] + [(chr(c), c + KC_A - ord('a')) for c in range(ord('a'), 50 ord('z') + 1)]) # Characters a-z. 51 52 53 def parse_file(file_name: str) -> List[Tuple[str, str]]: 54 """Parses autocorrections dictionary file. 55 Each line of the file defines one typo and its correction with the syntax 56 "typo -> correction". Blank lines or lines starting with '#' are ignored. The 57 function validates that typos only have characters a-z and that typos are not 58 substrings of other typos, otherwise the longer typo would never trigger. 59 Args: 60 file_name: String, path of the autocorrections dictionary. 61 Returns: 62 List of (typo, correction) tuples. 63 """ 64 65 try: 66 import english_words 67 correct_words = english_words.get_english_words_set(['web2'], lower=True, alpha=True) 68 except AttributeError: 69 from english_words import english_words_lower_alpha_set as correct_words 70 if not cli.args.quiet: 71 cli.echo('The english_words package is outdated, update by running:') 72 cli.echo(' {fg_cyan}python3 -m pip install english_words --upgrade') 73 except ImportError: 74 if not cli.args.quiet: 75 cli.echo('Autocorrection will falsely trigger when a typo is a substring of a correctly spelled word.') 76 cli.echo('To check for this, install the english_words package and rerun this script:') 77 cli.echo(' {fg_cyan}python3 -m pip install english_words') 78 # Use a minimal word list as a fallback. 79 correct_words = ('information', 'available', 'international', 'language', 'loosest', 'reference', 'wealthier', 'entertainment', 'association', 'provides', 'technology', 'statehood') 80 81 autocorrections = [] 82 typos = set() 83 for line_number, typo, correction in parse_file_lines(file_name): 84 if typo in typos: 85 cli.log.warning('{fg_red}Error:%d:{fg_reset} Ignoring duplicate typo: "{fg_cyan}%s{fg_reset}"', line_number, typo) 86 continue 87 88 # Check that `typo` is valid. 89 if not (all([c in TYPO_CHARS for c in typo])): 90 cli.log.error('{fg_red}Error:%d:{fg_reset} Typo "{fg_cyan}%s{fg_reset}" has characters other than a-z, \' and :.', line_number, typo) 91 maybe_exit(1) 92 for other_typo in typos: 93 if typo in other_typo or other_typo in typo: 94 cli.log.error('{fg_red}Error:%d:{fg_reset} Typos may not be substrings of one another, otherwise the longer typo would never trigger: "{fg_cyan}%s{fg_reset}" vs. "{fg_cyan}%s{fg_reset}".', line_number, typo, other_typo) 95 maybe_exit(1) 96 if len(typo) < 5: 97 cli.log.warning('{fg_yellow}Warning:%d:{fg_reset} It is suggested that typos are at least 5 characters long to avoid false triggers: "{fg_cyan}%s{fg_reset}"', line_number, typo) 98 if len(typo) > 127: 99 cli.log.error('{fg_red}Error:%d:{fg_reset} Typo exceeds 127 chars: "{fg_cyan}%s{fg_reset}"', line_number, typo) 100 maybe_exit(1) 101 102 check_typo_against_dictionary(typo, line_number, correct_words) 103 104 autocorrections.append((typo, correction)) 105 typos.add(typo) 106 107 return autocorrections 108 109 110 def make_trie(autocorrections: List[Tuple[str, str]]) -> Dict[str, Any]: 111 """Makes a trie from the the typos, writing in reverse. 112 Args: 113 autocorrections: List of (typo, correction) tuples. 114 Returns: 115 Dict of dict, representing the trie. 116 """ 117 trie = {} 118 for typo, correction in autocorrections: 119 node = trie 120 for letter in typo[::-1]: 121 node = node.setdefault(letter, {}) 122 node['LEAF'] = (typo, correction) 123 124 return trie 125 126 127 def parse_file_lines(file_name: str) -> Iterator[Tuple[int, str, str]]: 128 """Parses lines read from `file_name` into typo-correction pairs.""" 129 130 line_number = 0 131 for line in open(file_name, 'rt'): 132 line_number += 1 133 line = line.strip() 134 if line and line[0] != '#': 135 # Parse syntax "typo -> correction", using strip to ignore indenting. 136 tokens = [token.strip() for token in line.split('->', 1)] 137 if len(tokens) != 2 or not tokens[0]: 138 print(f'Error:{line_number}: Invalid syntax: "{line}"') 139 maybe_exit(1) 140 141 typo, correction = tokens 142 typo = typo.lower() # Force typos to lowercase. 143 typo = typo.replace(' ', ':') 144 145 yield line_number, typo, correction 146 147 148 def check_typo_against_dictionary(typo: str, line_number: int, correct_words) -> None: 149 """Checks `typo` against English dictionary words.""" 150 151 if typo.startswith(':') and typo.endswith(':'): 152 if typo[1:-1] in correct_words: 153 cli.log.warning('{fg_yellow}Warning:%d:{fg_reset} Typo "{fg_cyan}%s{fg_reset}" is a correctly spelled dictionary word.', line_number, typo) 154 elif typo.startswith(':') and not typo.endswith(':'): 155 for word in correct_words: 156 if word.startswith(typo[1:]): 157 cli.log.warning('{fg_yellow}Warning:%d: {fg_reset}Typo "{fg_cyan}%s{fg_reset}" would falsely trigger on correctly spelled word "{fg_cyan}%s{fg_reset}".', line_number, typo, word) 158 elif not typo.startswith(':') and typo.endswith(':'): 159 for word in correct_words: 160 if word.endswith(typo[:-1]): 161 cli.log.warning('{fg_yellow}Warning:%d:{fg_reset} Typo "{fg_cyan}%s{fg_reset}" would falsely trigger on correctly spelled word "{fg_cyan}%s{fg_reset}".', line_number, typo, word) 162 elif not typo.startswith(':') and not typo.endswith(':'): 163 for word in correct_words: 164 if typo in word: 165 cli.log.warning('{fg_yellow}Warning:%d:{fg_reset} Typo "{fg_cyan}%s{fg_reset}" would falsely trigger on correctly spelled word "{fg_cyan}%s{fg_reset}".', line_number, typo, word) 166 167 168 def serialize_trie(autocorrections: List[Tuple[str, str]], trie: Dict[str, Any]) -> List[int]: 169 """Serializes trie and correction data in a form readable by the C code. 170 Args: 171 autocorrections: List of (typo, correction) tuples. 172 trie: Dict of dicts. 173 Returns: 174 List of ints in the range 0-255. 175 """ 176 table = [] 177 178 # Traverse trie in depth first order. 179 def traverse(trie_node): 180 if 'LEAF' in trie_node: # Handle a leaf trie node. 181 typo, correction = trie_node['LEAF'] 182 word_boundary_ending = typo[-1] == ':' 183 typo = typo.strip(':') 184 i = 0 # Make the autocorrection data for this entry and serialize it. 185 while i < min(len(typo), len(correction)) and typo[i] == correction[i]: 186 i += 1 187 backspaces = len(typo) - i - 1 + word_boundary_ending 188 assert 0 <= backspaces <= 63 189 correction = correction[i:] 190 bs_count = [backspaces + 128] 191 data = bs_count + list(bytes(correction, 'ascii')) + [0] 192 193 entry = {'data': data, 'links': [], 'byte_offset': 0} 194 table.append(entry) 195 elif len(trie_node) == 1: # Handle trie node with a single child. 196 c, trie_node = next(iter(trie_node.items())) 197 entry = {'chars': c, 'byte_offset': 0} 198 199 # It's common for a trie to have long chains of single-child nodes. We 200 # find the whole chain so that we can serialize it more efficiently. 201 while len(trie_node) == 1 and 'LEAF' not in trie_node: 202 c, trie_node = next(iter(trie_node.items())) 203 entry['chars'] += c 204 205 table.append(entry) 206 entry['links'] = [traverse(trie_node)] 207 else: # Handle trie node with multiple children. 208 entry = {'chars': ''.join(sorted(trie_node.keys())), 'byte_offset': 0} 209 table.append(entry) 210 entry['links'] = [traverse(trie_node[c]) for c in entry['chars']] 211 return entry 212 213 traverse(trie) 214 215 def serialize(e: Dict[str, Any]) -> List[int]: 216 if not e['links']: # Handle a leaf table entry. 217 return e['data'] 218 elif len(e['links']) == 1: # Handle a chain table entry. 219 return [TYPO_CHARS[c] for c in e['chars']] + [0] # + encode_link(e['links'][0])) 220 else: # Handle a branch table entry. 221 data = [] 222 for c, link in zip(e['chars'], e['links']): 223 data += [TYPO_CHARS[c] | (0 if data else 64)] + encode_link(link) 224 return data + [0] 225 226 byte_offset = 0 227 for e in table: # To encode links, first compute byte offset of each entry. 228 e['byte_offset'] = byte_offset 229 byte_offset += len(serialize(e)) 230 assert 0 <= byte_offset <= 0xffff 231 232 return [b for e in table for b in serialize(e)] # Serialize final table. 233 234 235 def encode_link(link: Dict[str, Any]) -> List[int]: 236 """Encodes a node link as two bytes.""" 237 byte_offset = link['byte_offset'] 238 if not (0 <= byte_offset <= 0xffff): 239 cli.log.error('{fg_red}Error:{fg_reset} The autocorrection table is too large, a node link exceeds 64KB limit. Try reducing the autocorrection dict to fewer entries.') 240 maybe_exit(1) 241 return [byte_offset & 255, byte_offset >> 8] 242 243 244 def typo_len(e: Tuple[str, str]) -> int: 245 return len(e[0]) 246 247 248 def to_hex(b: int) -> str: 249 return f'0x{b:02X}' 250 251 252 @cli.argument('filename', type=normpath, help='The autocorrection database file') 253 @cli.argument('-kb', '--keyboard', type=keyboard_folder, completer=keyboard_completer, help='The keyboard to build a firmware for. Ignored when a output file is supplied.') 254 @cli.argument('-km', '--keymap', completer=keymap_completer, help='The keymap to build a firmware for. Ignored when a output file is supplied.') 255 @cli.argument('-o', '--output', arg_only=True, type=normpath, help='File to write to') 256 @cli.argument('-q', '--quiet', arg_only=True, action='store_true', help="Quiet mode, only output error messages") 257 @cli.subcommand('Generate the autocorrection data file from a dictionary file.') 258 def generate_autocorrect_data(cli): 259 autocorrections = parse_file(cli.args.filename) 260 trie = make_trie(autocorrections) 261 data = serialize_trie(autocorrections, trie) 262 263 current_keyboard = cli.args.keyboard or cli.config.user.keyboard or cli.config.generate_autocorrect_data.keyboard 264 current_keymap = cli.args.keymap or cli.config.user.keymap or cli.config.generate_autocorrect_data.keymap 265 266 if not cli.args.output and current_keyboard and current_keymap: 267 cli.args.output = locate_keymap(current_keyboard, current_keymap).parent / 'autocorrect_data.h' 268 269 assert all(0 <= b <= 255 for b in data) 270 271 min_typo = min(autocorrections, key=typo_len)[0] 272 max_typo = max(autocorrections, key=typo_len)[0] 273 274 # Build the autocorrect_data.h file. 275 autocorrect_data_h_lines = [GPL2_HEADER_C_LIKE, GENERATED_HEADER_C_LIKE, '#pragma once', ''] 276 277 autocorrect_data_h_lines.append(f'// Autocorrection dictionary ({len(autocorrections)} entries):') 278 for typo, correction in autocorrections: 279 autocorrect_data_h_lines.append(f'// {typo:<{len(max_typo)}} -> {correction}') 280 281 autocorrect_data_h_lines.append('') 282 autocorrect_data_h_lines.append(f'#define AUTOCORRECT_MIN_LENGTH {len(min_typo)} // "{min_typo}"') 283 autocorrect_data_h_lines.append(f'#define AUTOCORRECT_MAX_LENGTH {len(max_typo)} // "{max_typo}"') 284 autocorrect_data_h_lines.append(f'#define DICTIONARY_SIZE {len(data)}') 285 autocorrect_data_h_lines.append('') 286 autocorrect_data_h_lines.append('static const uint8_t autocorrect_data[DICTIONARY_SIZE] PROGMEM = {') 287 autocorrect_data_h_lines.append(textwrap.fill(' %s' % (', '.join(map(to_hex, data))), width=100, subsequent_indent=' ')) 288 autocorrect_data_h_lines.append('};') 289 290 # Show the results 291 dump_lines(cli.args.output, autocorrect_data_h_lines, cli.args.quiet)