hstore.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. import json
  2. from django.contrib.postgres import forms, lookups
  3. from django.contrib.postgres.fields.array import ArrayField
  4. from django.core import exceptions
  5. from django.db.models import Field, TextField, Transform
  6. from django.db.models.fields.mixins import CheckFieldDefaultMixin
  7. from django.utils.translation import gettext_lazy as _
  8. __all__ = ["HStoreField"]
  9. class HStoreField(CheckFieldDefaultMixin, Field):
  10. empty_strings_allowed = False
  11. description = _("Map of strings to strings/nulls")
  12. default_error_messages = {
  13. "not_a_string": _("The value of “%(key)s” is not a string or null."),
  14. }
  15. _default_hint = ("dict", "{}")
  16. def db_type(self, connection):
  17. return "hstore"
  18. def get_transform(self, name):
  19. transform = super().get_transform(name)
  20. if transform:
  21. return transform
  22. return KeyTransformFactory(name)
  23. def validate(self, value, model_instance):
  24. super().validate(value, model_instance)
  25. for key, val in value.items():
  26. if not isinstance(val, str) and val is not None:
  27. raise exceptions.ValidationError(
  28. self.error_messages["not_a_string"],
  29. code="not_a_string",
  30. params={"key": key},
  31. )
  32. def to_python(self, value):
  33. if isinstance(value, str):
  34. value = json.loads(value)
  35. return value
  36. def value_to_string(self, obj):
  37. return json.dumps(self.value_from_object(obj))
  38. def formfield(self, **kwargs):
  39. return super().formfield(
  40. **{
  41. "form_class": forms.HStoreField,
  42. **kwargs,
  43. }
  44. )
  45. def get_prep_value(self, value):
  46. value = super().get_prep_value(value)
  47. if isinstance(value, dict):
  48. prep_value = {}
  49. for key, val in value.items():
  50. key = str(key)
  51. if val is not None:
  52. val = str(val)
  53. prep_value[key] = val
  54. value = prep_value
  55. if isinstance(value, list):
  56. value = [str(item) for item in value]
  57. return value
  58. HStoreField.register_lookup(lookups.DataContains)
  59. HStoreField.register_lookup(lookups.ContainedBy)
  60. HStoreField.register_lookup(lookups.HasKey)
  61. HStoreField.register_lookup(lookups.HasKeys)
  62. HStoreField.register_lookup(lookups.HasAnyKeys)
  63. class KeyTransform(Transform):
  64. output_field = TextField()
  65. def __init__(self, key_name, *args, **kwargs):
  66. super().__init__(*args, **kwargs)
  67. self.key_name = key_name
  68. def as_sql(self, compiler, connection):
  69. lhs, params = compiler.compile(self.lhs)
  70. return "(%s -> %%s)" % lhs, tuple(params) + (self.key_name,)
  71. class KeyTransformFactory:
  72. def __init__(self, key_name):
  73. self.key_name = key_name
  74. def __call__(self, *args, **kwargs):
  75. return KeyTransform(self.key_name, *args, **kwargs)
  76. @HStoreField.register_lookup
  77. class KeysTransform(Transform):
  78. lookup_name = "keys"
  79. function = "akeys"
  80. output_field = ArrayField(TextField())
  81. @HStoreField.register_lookup
  82. class ValuesTransform(Transform):
  83. lookup_name = "values"
  84. function = "avals"
  85. output_field = ArrayField(TextField())