manager.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. from IPy import IP
  2. from django.core.exceptions import ValidationError
  3. from django.db import models, connection
  4. from django.db.models import sql, query
  5. INET_TERMS = {
  6. 'inet_lt': '<',
  7. 'inet_lte': '<=',
  8. 'inet_exact': '=',
  9. 'inet_gte': '>=',
  10. 'inet_gt': '>',
  11. 'inet_not': '<>',
  12. 'inet_is_contained': '<<',
  13. 'inet_is_contained_or_equal': '<<=',
  14. 'inet_contains': '>>',
  15. 'inet_contains': '>>=',
  16. }
  17. class IPQuery(sql.Query):
  18. query_terms = sql.Query.query_terms.copy()
  19. query_terms.update(INET_TERMS)
  20. def add_filter(self, (filter_string, value), *args, **kwargs):
  21. if isinstance(value, IP):
  22. value = unicode(value)
  23. return super(IPQuery, self).add_filter((filter_string, value), *args, **kwargs)
  24. class IPWhere(sql.where.WhereNode):
  25. def make_atom(self, child, qn):
  26. table_alias, name, db_type, lookup_type, value_annot, params = child
  27. if lookup_type in INET_TERMS:
  28. return ('%s.%s %s inet %%s' % (table_alias, name, INET_TERMS[lookup_type]), params)
  29. return super(IPWhere, self).make_atom(child, qn)
  30. class IPManger(models.Manager):
  31. def get_query_set(self):
  32. q = IPQuery(self.model, connection, IPWhere)
  33. return query.QuerySet(self.model, q)
  34. class IPAddressField(models.IPAddressField):
  35. __metaclass__ = models.SubfieldBase
  36. def db_type(self):
  37. return 'inet'
  38. def to_python(self, value):
  39. if not value:
  40. return None
  41. try:
  42. return IP(value)
  43. except ValueError, e:
  44. raise ValidationError(e)
  45. def get_db_prep_value(self, value):
  46. return unicode(self.to_python(value))
  47. def get_db_prep_lookup(self, lookup_type, value):
  48. value = unicode(value)
  49. if lookup_type in INET_TERMS:
  50. return [value]
  51. return super(IPAddressField, self).get_db_prep_lookup(lookup_type, value)
  52. class Foo(models.Model):
  53. ip = IPAddressField()
  54. objects = IPManger()
  55. def __unicode__(self):
  56. return unicode(self.ip)