Implement finding documents by tags

1 files changed, 43 insertions, 0 deletions

diff --git a/archivist/cli.py b/archivist/cli.py
index 4fd78e4..5085ab7 100644
--- a/archivist/cli.py
+++ b/archivist/cli.py
@@ -1,4 +1,6 @@
 import click
+from functools import reduce
+import operator as op
 
 from .prefixes import query_pseudo_prefix, is_pseudo_prefix
 
@@ -111,6 +113,13 @@ class PrefixTag:
     def prefixed_name(self):
         return prefix_tag_name(self.tag, self.prefix)
 
+    def is_virtual(self):
+        return self.prefix and is_pseudo_prefix(self.prefix)
+
+    @property
+    def virtual_query(self):
+        return query_pseudo_prefix(self.prefix, self.tag)
+
     def __str__(self):
         return self.prefixed_name()
 
@@ -247,3 +256,37 @@ def add_doc(file, tags, create_tags, ignore_missing_tags):
 
         for t in tags:
             DocumentTag.create(document = doc, tag = t)
+
+@doc.command('find')
+@click.argument('tags', type=TAG, nargs=-1)
+def find_doc(tags):
+    from .model import Document, DocumentTag, TagClosure, Tag
+
+    pseudo_tags = []
+    normal_tags = []
+
+    for t in tags:
+        if t.is_virtual():
+            pseudo_tags.append(t)
+        else:
+            normal_tags.append(t)
+
+    query = Document.select()
+
+    if normal_tags:
+        tag_query = None
+        for t in fetch_tags(normal_tags):
+            desc = TagClosure.descendants(t, include_node=True).select(Tag.id)
+            subq = DocumentTag.select(DocumentTag.document_id).where(DocumentTag.tag << desc)
+            if tag_query is None:
+                tag_query = subq
+            else:
+                tag_query = tag_query & subq
+
+        query = query.where(Document.id << tag_query)
+
+    if pseudo_tags:
+        query = query.where(reduce(op.and_, (p.virtual_query for p in pseudo_tags)))
+
+    for doc in query.iterator():
+        print("* ID %d -- %s" % (doc.id, doc.original_path))