DocType subclasses can now be used directly to limit Searches

ref #60

docs/persistence.rst

@@ -131,14 +131,13 @@ To search for this document type, use the ``search`` class method:
     for posts in results:
         print(post._meta.score, post.title)
 
-Alternatively you can just take a ``Search`` object and register the
-``from_es`` method of you document class as a callback for certain document
-types:
+Alternatively you can just take a ``Search`` object and restrict it to return
+our document type, wrapped in correct class:
 
 .. code:: python
 
     s = Search()
-    s = s.doc_type(post=Post.from_es)
+    s = s.doc_type(Post)
 
 You can also combine document classes with standard doc types (just strings),
 which will be treated as before. You can also pass in multiple ``DocType``

elasticsearch_dsl/search.py

@@ -80,12 +80,13 @@ def __init__(self, using='default', index=None, doc_type=None, extra=None):
         self._doc_type = []
         self._doc_type_map = {}
         if isinstance(doc_type, (tuple, list)):
-            self._doc_type.extend(doc_type)
+            for dt in doc_type:
+                self._add_doc_type(dt)
         elif isinstance(doc_type, dict):
             self._doc_type.extend(doc_type.keys())
             self._doc_type_map.update(doc_type)
         elif doc_type:
-            self._doc_type.append(doc_type)
+            self._add_doc_type(doc_type)
 
         self.query = ProxyQuery(self, 'query')
         self.filter = ProxyFilter(self, 'filter')
@@ -302,10 +303,16 @@ def index(self, *index):
             s._index = (self._index or []) + list(index)
         return s
 
+    def _add_doc_type(self, doc_type):
+        if hasattr(doc_type, '_doc_type'):
+            self._doc_type_map[doc_type._doc_type.name] = doc_type.from_es
+            doc_type = doc_type._doc_type.name
+        self._doc_type.append(doc_type)
+
     def doc_type(self, *doc_type, **kwargs):
         """
         Set the type to search through. You can supply a single value or
-        multiple.
+        multiple. Values can be strings or subclasses of ``DocType``.
         
         You can also pass in any keyword arguments, mapping a doc_type to a
         callback that should be used instead of the Result class.
@@ -315,15 +322,16 @@ def doc_type(self, *doc_type, **kwargs):
 
         Example:
 
-            s = Search().doc_type('product', 'store', user=User.from_es)
+            s = Search().doc_type('product', 'store', User, custom=my_callback)
         """
         # .doc_type() resets
         s = self._clone()
         if not doc_type and not kwargs:
             s._doc_type = []
             s._doc_type_map = {}
         else:
-            s._doc_type.extend(doc_type)
+            for dt in doc_type:
+                s._add_doc_type(dt)
             s._doc_type.extend(kwargs.keys())
             s._doc_type_map.update(kwargs)
         return s

test_elasticsearch_dsl/test_search.py

@@ -1,7 +1,6 @@
 from copy import deepcopy
 
-from elasticsearch_dsl import search, query, F, Q
-from elasticsearch_dsl.connections import connections
+from elasticsearch_dsl import search, query, F, Q, DocType
 
 def test_search_starts_with_empty_query():
     s = search.Search()
@@ -103,6 +102,20 @@ def test_search_doc_type():
     assert s._doc_type == ['i', 'i2']
     assert s2._doc_type == ['i', 'i2', 'i3']
 
+
+def test_doc_type_can_be_document_class():
+    class MyDocType(DocType):
+        pass
+
+    s = search.Search(doc_type=MyDocType)
+    assert s._doc_type == ['my_doc_type']
+    assert s._doc_type_map == {'my_doc_type': MyDocType.from_es}
+
+    s = search.Search().doc_type(MyDocType)
+    assert s._doc_type == ['my_doc_type']
+    assert s._doc_type_map == {'my_doc_type': MyDocType.from_es}
+
+
 def test_sort():
     s = search.Search()
     s = s.sort('fielda', '-fieldb')