]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/tokenizer/icu_name_processor.py
adapt tests for ICU tokenizer
[nominatim.git] / nominatim / tokenizer / icu_name_processor.py
1 """
2 Processor for names that are imported into the database based on the
3 ICU library.
4 """
5 import json
6 import itertools
7
8 from icu import Transliterator
9 import datrie
10
11 from nominatim.db.properties import set_property, get_property
12
13 DBCFG_IMPORT_NORM_RULES = "tokenizer_import_normalisation"
14 DBCFG_IMPORT_TRANS_RULES = "tokenizer_import_transliteration"
15 DBCFG_IMPORT_REPLACEMENTS = "tokenizer_import_replacements"
16 DBCFG_SEARCH_STD_RULES = "tokenizer_search_standardization"
17
18
19 class ICUNameProcessorRules:
20     """ Data object that saves the rules needed for the name processor.
21
22         The rules can either be initialised through an ICURuleLoader or
23         be loaded from a database when a connection is given.
24     """
25     def __init__(self, loader=None, conn=None):
26         if loader is not None:
27             self.norm_rules = loader.get_normalization_rules()
28             self.trans_rules = loader.get_transliteration_rules()
29             self.replacements = loader.get_replacement_pairs()
30             self.search_rules = loader.get_search_rules()
31         elif conn is not None:
32             self.norm_rules = get_property(conn, DBCFG_IMPORT_NORM_RULES)
33             self.trans_rules = get_property(conn, DBCFG_IMPORT_TRANS_RULES)
34             self.replacements = json.loads(get_property(conn, DBCFG_IMPORT_REPLACEMENTS))
35             self.search_rules = get_property(conn, DBCFG_SEARCH_STD_RULES)
36         else:
37             assert False, "Parameter loader or conn required."
38
39         # Compute the set of characters used in the replacement list.
40         # We need this later when computing the tree.
41         chars = set()
42         for full, repl in self.replacements:
43             chars.update(full)
44             for word in repl:
45                 chars.update(word)
46         self.replacement_charset = ''.join(chars)
47
48
49     def save_rules(self, conn):
50         """ Save the rules in the property table of the given database.
51             the rules can be loaded again by handing in a connection into
52             the constructor of the class.
53         """
54         set_property(conn, DBCFG_IMPORT_NORM_RULES, self.norm_rules)
55         set_property(conn, DBCFG_IMPORT_TRANS_RULES, self.trans_rules)
56         set_property(conn, DBCFG_IMPORT_REPLACEMENTS, json.dumps(self.replacements))
57         set_property(conn, DBCFG_SEARCH_STD_RULES, self.search_rules)
58
59
60 class ICUNameProcessor:
61
62     def __init__(self, rules):
63         self.normalizer = Transliterator.createFromRules("icu_normalization",
64                                                          rules.norm_rules)
65         self.to_ascii = Transliterator.createFromRules("icu_to_ascii",
66                                                        rules.trans_rules)
67         self.search = Transliterator.createFromRules("icu_search",
68                                                      rules.search_rules)
69
70         self.replacements = datrie.Trie(rules.replacement_charset)
71         for full, repl in rules.replacements:
72             self.replacements[full] = repl
73
74
75     def get_normalized(self, name):
76         """ Normalize the given name, i.e. remove all elements not relevant
77             for search.
78         """
79         return self.normalizer.transliterate(name).strip()
80
81     def get_variants_ascii(self, norm_name):
82         """ Compute the spelling variants for the given normalized name
83             and transliterate the result.
84         """
85         baseform = ' ' + norm_name + ' '
86         variants = ['']
87
88         startpos = 0
89         pos = 0
90         while pos < len(baseform):
91             full, repl = self.replacements.longest_prefix_item(baseform[pos:],
92                                                                (None, None))
93             if full is not None:
94                 done = baseform[startpos:pos]
95                 variants = [v + done + r for v, r in itertools.product(variants, repl)]
96                 startpos = pos + len(full)
97                 pos = startpos
98             else:
99                 pos += 1
100
101         if startpos == 0:
102             return [self.to_ascii.transliterate(norm_name)]
103
104         return [self.to_ascii.transliterate(v + baseform[startpos:pos]).strip() for v in variants]
105
106
107     def get_search_normalized(self, name):
108         """ Return the normalized version of the name (including transliteration)
109             to be applied at search time.
110         """
111         return self.search.transliterate(' ' + name + ' ').strip()