nominatim/tokenizer/icu_variants.py

   1 """
   2 Data structures for saving variant expansions for ICU tokenizer.
   3 """
   4 from collections import namedtuple
   5 import json
   6
   7 from nominatim.errors import UsageError
   8
   9 _ICU_VARIANT_PORPERTY_FIELDS = ['lang']
  10
  11 def _get_strtuple_prop(rules, field):
  12     """ Return the given field of the rules dictionary as a list.
  13
  14         If the field is not defined or empty, returns None. If the field is
  15         a singe string, it is converted into a tuple with a single element.
  16         If the field is a list of strings, return as a string tuple.
  17         Raise a usage error in all other cases.
  18     """
  19     value = rules.get(field)
  20
  21     if not value:
  22         return None
  23
  24     if isinstance(value, str):
  25         return (value,)
  26
  27     if not isinstance(value, list) or any(not isinstance(x, str) for x in value):
  28         raise UsageError("YAML variant property '{}' should be a list.".format(field))
  29
  30     return tuple(value)
  31
  32
  33 class ICUVariantProperties(namedtuple('_ICUVariantProperties', _ICU_VARIANT_PORPERTY_FIELDS,
  34                                       defaults=(None, )*len(_ICU_VARIANT_PORPERTY_FIELDS))):
  35     """ Data container for saving properties that describe when a variant
  36         should be applied.
  37
  38         Porperty instances are hashable.
  39     """
  40     @classmethod
  41     def from_rules(cls, rules):
  42         """ Create a new property type from a generic dictionary.
  43
  44             The function only takes into account the properties that are
  45             understood presently and ignores all others.
  46         """
  47         return cls(lang=_get_strtuple_prop(rules, 'lang'))
  48
  49
  50 ICUVariant = namedtuple('ICUVariant', ['source', 'replacement', 'properties'])
  51
  52 def pickle_variant_set(variants):
  53     """ Serializes an iterable of variant rules to a string.
  54     """
  55     # Create a list of property sets. So they don't need to be duplicated
  56     properties = {}
  57     pid = 1
  58     for variant in variants:
  59         if variant.properties not in properties:
  60             properties[variant.properties] = pid
  61             pid += 1
  62
  63     # Convert the variants into a simple list.
  64     variants = [(v.source, v.replacement, properties[v.properties]) for v in variants]
  65
  66     # Convert everythin to json.
  67     return json.dumps({'properties': {v: k._asdict() for k, v in properties.items()},
  68                        'variants': variants})
  69
  70
  71 def unpickle_variant_set(variant_string):
  72     """ Deserializes a variant string that was previously created with
  73         pickle_variant_set() into a set of ICUVariants.
  74     """
  75     data = json.loads(variant_string)
  76
  77     properties = {int(k): ICUVariantProperties(**v) for k, v in data['properties'].items()}
  78     print(properties)
  79
  80     return set((ICUVariant(src, repl, properties[pid]) for src, repl, pid in data['variants']))