[WIP] Added support for declarative document types.

elasticutils/__init__.py

@@ -1,3 +1,4 @@
+import collections
 import copy
 import logging
 from datetime import datetime
@@ -7,6 +8,7 @@
 from elasticsearch.helpers import bulk_index
 
 from elasticutils._version import __version__  # noqa
+from elasticutils.fields import SearchField
 
 
 log = logging.getLogger('elasticutils')
@@ -2194,3 +2196,33 @@ def refresh_index(cls, es=None, index=None):
             index = cls.get_index()
 
         es.indices.refresh(index=index)
+
+
+class DeclarativeMappingMeta(type):
+
+    def __new__(cls, name, bases, attrs):
+        fields = [(name_, attrs.pop(name_)) for name_, column in attrs.items()
+                  if isinstance(column, SearchField)]
+        # Put fields in order defined in the class.
+        fields.sort(key=lambda f: f[1]._creation_order)
+        attrs['fields'] = fields
+        return super(DeclarativeMappingMeta, cls).__new__(cls, name, bases,
+                                                          attrs)
+
+
+class DocumentType(object):
+    __metaclass__ = DeclarativeMappingMeta
+
+    def get_mapping(self):
+        """
+        Returns mapping based on defined fields.
+        """
+        fields = collections.OrderedDict()
+        for name, field in self.fields:
+            name = field.index_fieldname or name
+            defn = field.get_definition()
+            fields[name] = defn
+
+        mapping = {'properties': fields}
+
+        return mapping

elasticutils/exceptions.py

@@ -0,0 +1,3 @@
+class SearchFieldError(Exception):
+    """Raised when a field encounters an error."""
+    pass

elasticutils/fields.py

@@ -0,0 +1,222 @@
+import base64
+import datetime
+import re
+from decimal import Decimal
+
+from .exceptions import SearchFieldError
+
+
+DATE_REGEX = re.compile('^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2}).*?$')
+DATETIME_REGEX = re.compile('^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})'
+                            '(T|\s+)(?P<hour>\d{2}):(?P<minute>\d{2}):'
+                            '(?P<second>\d{2}).*?$')
+
+
+class SearchField(object):
+
+    field_type = None
+    attrs = []
+
+    # Used to maintain the order of fields as defined in the class.
+    _creation_order = 0
+
+    def __init__(self, *args, **kwargs):
+        # These are special.
+        for attr in ('index_fieldname', 'is_multivalue'):
+            setattr(self, attr, kwargs.pop(attr, None))
+
+        # Set all kwargs on self for later access.
+        for attr in kwargs.keys():
+            self.attrs.append(attr)
+            setattr(self, attr, kwargs.pop(attr, None))
+
+        # Store this fields order.
+        self._creation_order = SearchField._creation_order
+        # Increment order number for future fields.
+        SearchField._creation_order += 1
+
+    def to_es(self, value):
+        """
+        Converts a Python value to an Elasticsearch value.
+
+        Extending classes should override this method.
+        """
+        return value
+
+    def to_python(self, value):
+        """
+        Converts an Elasticsearch value to a Python value.
+
+        Extending classes should override this method.
+        """
+        return value
+
+    def get_definition(self):
+        """
+        Returns the resprentation for this field's definition in the mapping.
+        """
+        f = {'type': self.field_type}
+
+        for attr in self.attrs:
+            val = getattr(self, attr, None)
+            if val is not None:
+                f[attr] = val
+
+        return f
+
+
+class StringField(SearchField):
+    field_type = 'string'
+
+    def to_es(self, value):
+        if value is None:
+            return None
+
+        return unicode(value)
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        return unicode(value)
+
+
+class IntegerField(SearchField):
+    field_type = 'integer'
+
+    def __init__(self, type='integer', *args, **kwargs):
+        if type in ('byte', 'short', 'integer', 'long'):
+            self.field_type = type
+        super(IntegerField, self).__init__(*args, **kwargs)
+
+    def to_es(self, value):
+        if value is None:
+            return None
+
+        return int(value)
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        return int(value)
+
+
+class FloatField(SearchField):
+    field_type = 'float'
+
+    def __init__(self, type='float', *args, **kwargs):
+        if type in ('float', 'double'):
+            self.field_type = type
+        super(FloatField, self).__init__(*args, **kwargs)
+
+    def to_es(self, value):
+        if value is None:
+            return None
+
+        return float(value)
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        return float(value)
+
+
+class DecimalField(StringField):
+
+    def to_es(self, value):
+        if value is None:
+            return None
+
+        return str(float(value))
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        return Decimal(str(value))
+
+
+class BooleanField(SearchField):
+    field_type = 'boolean'
+
+    def to_es(self, value):
+        if value is None:
+            return None
+
+        return bool(value)
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        return bool(value)
+
+
+class DateField(SearchField):
+    field_type = 'date'
+
+    def to_es(self, value):
+        if isinstance(value, (datetime.date, datetime.datetime)):
+            return value.isoformat()
+
+        return value
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        if isinstance(value, basestring):
+            match = DATE_REGEX.search(value)
+
+            if match:
+                data = match.groupdict()
+                return datetime.date(
+                    int(data['year']), int(data['month']), int(data['day']))
+            else:
+                raise SearchFieldError(
+                    "Date provided to '%s' field doesn't appear to be a valid "
+                    "date string: '%s'" % (self.instance_name, value))
+
+        return value
+
+
+class DateTimeField(DateField):
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        if isinstance(value, basestring):
+            match = DATETIME_REGEX.search(value)
+
+            if match:
+                data = match.groupdict()
+                return datetime.datetime(
+                    int(data['year']), int(data['month']), int(data['day']),
+                    int(data['hour']), int(data['minute']),
+                    int(data['second']))
+            else:
+                raise SearchFieldError(
+                    "Datetime provided to '%s' field doesn't appear to be a "
+                    "valid datetime string: '%s'" % (
+                        self.instance_name, value))
+
+        return value
+
+
+class BinaryField(SearchField):
+    field_type = 'binary'
+
+    def to_es(self, value):
+        if value is None:
+            return None
+
+        return base64.b64encode(value)
+
+    def to_python(self, value):
+        if value is None:
+            return None
+
+        return base64.b64decode(value)

elasticutils/tests/test_document_type.py

@@ -0,0 +1,42 @@
+from unittest import TestCase
+
+from nose.tools import eq_
+
+from elasticutils import DocumentType, fields
+
+
+class BookDocumentType(DocumentType):
+
+    id = fields.IntegerField(type='long')
+    name = fields.StringField(analyzer='snowball')
+    name_sort = fields.StringField(index='not_analyzed')
+    authors = fields.StringField(is_multivalued=True)
+    published_date = fields.DateField()
+    price = fields.DecimalField()
+    is_autographed = fields.BooleanField()
+    sales = fields.IntegerField()
+
+
+class DocumentTypeTest(TestCase):
+
+    def setUp(self):
+        self._type = BookDocumentType
+
+    def test_mapping(self):
+        mapping = self._type().get_mapping()
+
+        # Check top level element.
+        eq_(mapping.keys(), ['properties'])
+
+        fields = mapping['properties']
+
+        eq_(fields['id']['type'], 'long')
+        eq_(fields['name']['type'], 'string')
+        eq_(fields['name']['analyzer'], 'snowball')
+        eq_(fields['name_sort']['type'], 'string')
+        eq_(fields['name_sort']['index'], 'not_analyzed')
+        eq_(fields['authors']['type'], 'string')
+        eq_(fields['published_date']['type'], 'date')
+        eq_(fields['price']['type'], 'string')
+        eq_(fields['is_autographed']['type'], 'boolean')
+        eq_(fields['sales']['type'], 'integer')