test/python/mock_icu_word_table.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2025 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Legacy word table for testing with functions to prefil and test contents
   9 of the table.
  10 """
  11 from nominatim_db.db.connection import execute_scalar
  12
  13
  14 class MockIcuWordTable:
  15     """ A word table for testing using legacy word table structure.
  16     """
  17     def __init__(self, conn):
  18         self.conn = conn
  19         with conn.cursor() as cur:
  20             cur.execute("""CREATE TABLE word (word_id INTEGER,
  21                                               word_token text NOT NULL,
  22                                               type text NOT NULL,
  23                                               word text,
  24                                               info jsonb)""")
  25
  26         conn.commit()
  27
  28     def add_full_word(self, word_id, word, word_token=None):
  29         with self.conn.cursor() as cur:
  30             cur.execute("""INSERT INTO word (word_id, word_token, type, word, info)
  31                            VALUES(%s, %s, 'W', %s, '{}'::jsonb)""",
  32                         (word_id, word or word_token, word))
  33         self.conn.commit()
  34
  35     def add_special(self, word_token, word, cls, typ, oper):
  36         with self.conn.cursor() as cur:
  37             cur.execute("""INSERT INTO word (word_token, type, word, info)
  38                               VALUES (%s, 'S', %s,
  39                                       json_build_object('class', %s::text,
  40                                                         'type', %s::text,
  41                                                         'op', %s::text))
  42                         """, (word_token, word, cls, typ, oper))
  43         self.conn.commit()
  44
  45     def add_country(self, country_code, word_token):
  46         with self.conn.cursor() as cur:
  47             cur.execute("""INSERT INTO word (word_token, type, word)
  48                            VALUES(%s, 'C', %s)""",
  49                         (word_token, country_code))
  50         self.conn.commit()
  51
  52     def add_postcode(self, word_token, postcode):
  53         with self.conn.cursor() as cur:
  54             cur.execute("""INSERT INTO word (word_token, type, word)
  55                               VALUES (%s, 'P', %s)
  56                         """, (word_token, postcode))
  57         self.conn.commit()
  58
  59     def add_housenumber(self, word_id, word_tokens, word=None):
  60         with self.conn.cursor() as cur:
  61             if isinstance(word_tokens, str):
  62                 # old style without analyser
  63                 cur.execute("""INSERT INTO word (word_id, word_token, type)
  64                                   VALUES (%s, %s, 'H')
  65                             """, (word_id, word_tokens))
  66             else:
  67                 if word is None:
  68                     word = word_tokens[0]
  69                 for token in word_tokens:
  70                     cur.execute("""INSERT INTO word (word_id, word_token, type, word, info)
  71                                       VALUES (%s, %s, 'H', %s,
  72                                               jsonb_build_object('lookup', %s::text))
  73                                 """, (word_id, token, word, word_tokens[0]))
  74
  75         self.conn.commit()
  76
  77     def count(self):
  78         return execute_scalar(self.conn, "SELECT count(*) FROM word")
  79
  80     def count_special(self):
  81         return execute_scalar(self.conn, "SELECT count(*) FROM word WHERE type = 'S'")
  82
  83     def count_housenumbers(self):
  84         return execute_scalar(self.conn, "SELECT count(*) FROM word WHERE type = 'H'")
  85
  86     def get_special(self):
  87         with self.conn.cursor() as cur:
  88             cur.execute("SELECT word_token, info, word FROM word WHERE type = 'S'")
  89             result = set(((row[0], row[2], row[1]['class'],
  90                            row[1]['type'], row[1]['op']) for row in cur))
  91             assert len(result) == cur.rowcount, "Word table has duplicates."
  92             return result
  93
  94     def get_country(self):
  95         with self.conn.cursor() as cur:
  96             cur.execute("SELECT word, word_token FROM word WHERE type = 'C'")
  97             result = set((tuple(row) for row in cur))
  98             assert len(result) == cur.rowcount, "Word table has duplicates."
  99             return result
 100
 101     def get_postcodes(self):
 102         with self.conn.cursor() as cur:
 103             cur.execute("SELECT word FROM word WHERE type = 'P'")
 104             return set((row[0] for row in cur))
 105
 106     def get_partial_words(self):
 107         with self.conn.cursor() as cur:
 108             cur.execute("SELECT word_token, info FROM word WHERE type ='w'")
 109             return set(((row[0], row[1]['count']) for row in cur))