fields.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. """
  2. Useful form fields for use with the Django ORM.
  3. """
  4. from __future__ import unicode_literals
  5. import datetime
  6. import operator
  7. try:
  8. from django.conf import settings
  9. from django.utils import timezone
  10. has_timezone = True
  11. except ImportError:
  12. has_timezone = False
  13. from wtforms import fields, widgets
  14. from wtforms.compat import string_types
  15. from wtforms.validators import ValidationError
  16. __all__ = (
  17. 'ModelSelectField', 'QuerySetSelectField', 'DateTimeField'
  18. )
  19. class QuerySetSelectField(fields.SelectFieldBase):
  20. """
  21. Given a QuerySet either at initialization or inside a view, will display a
  22. select drop-down field of choices. The `data` property actually will
  23. store/keep an ORM model instance, not the ID. Submitting a choice which is
  24. not in the queryset will result in a validation error.
  25. Specify `get_label` to customize the label associated with each option. If
  26. a string, this is the name of an attribute on the model object to use as
  27. the label text. If a one-argument callable, this callable will be passed
  28. model instance and expected to return the label text. Otherwise, the model
  29. object's `__str__` or `__unicode__` will be used.
  30. If `allow_blank` is set to `True`, then a blank choice will be added to the
  31. top of the list. Selecting this choice will result in the `data` property
  32. being `None`. The label for the blank choice can be set by specifying the
  33. `blank_text` parameter.
  34. """
  35. widget = widgets.Select()
  36. def __init__(self, label=None, validators=None, queryset=None, get_label=None, allow_blank=False, blank_text='', **kwargs):
  37. super(QuerySetSelectField, self).__init__(label, validators, **kwargs)
  38. self.allow_blank = allow_blank
  39. self.blank_text = blank_text
  40. self._set_data(None)
  41. if queryset is not None:
  42. self.queryset = queryset.all() # Make sure the queryset is fresh
  43. if get_label is None:
  44. self.get_label = lambda x: x
  45. elif isinstance(get_label, string_types):
  46. self.get_label = operator.attrgetter(get_label)
  47. else:
  48. self.get_label = get_label
  49. def _get_data(self):
  50. if self._formdata is not None:
  51. for obj in self.queryset:
  52. if obj.pk == self._formdata:
  53. self._set_data(obj)
  54. break
  55. return self._data
  56. def _set_data(self, data):
  57. self._data = data
  58. self._formdata = None
  59. data = property(_get_data, _set_data)
  60. def iter_choices(self):
  61. if self.allow_blank:
  62. yield ('__None', self.blank_text, self.data is None)
  63. for obj in self.queryset:
  64. yield (obj.pk, self.get_label(obj), obj == self.data)
  65. def process_formdata(self, valuelist):
  66. if valuelist:
  67. if valuelist[0] == '__None':
  68. self.data = None
  69. else:
  70. self._data = None
  71. self._formdata = int(valuelist[0])
  72. def pre_validate(self, form):
  73. if not self.allow_blank or self.data is not None:
  74. for obj in self.queryset:
  75. if self.data == obj:
  76. break
  77. else:
  78. raise ValidationError(self.gettext('Not a valid choice'))
  79. class ModelSelectField(QuerySetSelectField):
  80. """
  81. Like a QuerySetSelectField, except takes a model class instead of a
  82. queryset and lists everything in it.
  83. """
  84. def __init__(self, label=None, validators=None, model=None, **kwargs):
  85. super(ModelSelectField, self).__init__(label, validators, queryset=model._default_manager.all(), **kwargs)
  86. class DateTimeField(fields.DateTimeField):
  87. """
  88. Adds support for Django's timezone utilities.
  89. Requires Django >= 1.5
  90. """
  91. def __init__(self, *args, **kwargs):
  92. if not has_timezone:
  93. raise ImportError('DateTimeField requires Django >= 1.5')
  94. super(DateTimeField, self).__init__(*args, **kwargs)
  95. def process_formdata(self, valuelist):
  96. super(DateTimeField, self).process_formdata(valuelist)
  97. date = self.data
  98. if settings.USE_TZ and date is not None and timezone.is_naive(date):
  99. current_timezone = timezone.get_current_timezone()
  100. self.data = timezone.make_aware(date, current_timezone)
  101. def _value(self):
  102. date = self.data
  103. if settings.USE_TZ and isinstance(date, datetime.datetime) and timezone.is_aware(date):
  104. self.data = timezone.localtime(date)
  105. return super(DateTimeField, self)._value()