Allow to add implications for tags

1 files changed, 30 insertions, 12 deletions

diff --git a/archivist/cli.py b/archivist/cli.py
index 2e0fbfd..1ebe2db 100644
--- a/archivist/cli.py
+++ b/archivist/cli.py
@@ -111,6 +111,21 @@ def create_tag(name, description=None):
 
         return Tag.create_or_get(name = name.tag, prefix = name.prefix, description = description)
 
+def fetch_tags(tag_list, ignore_missing=False):
+    if not tag_list:
+        return []
+
+    from .model import Tag
+    from .peewee_ext import sqlite_tuple_in
+
+    fetched_tags = Tag.select().where(sqlite_tuple_in((Tag.prefix, Tag.name), tag_list)).execute()
+    if len(fetched_tags) < len(tag_list):
+        print("Some tags were not present:", ', '.join(set(map(str, tag_list)).difference(map(str, fetched_tags))))
+        if not ignore_missing:
+            raise click.ClickException("Not all tags present")
+
+    return fetched_tags
+
 @tag.command('add')
 @click.argument('name', type = TAG)
 @click.argument('description', required = False)
@@ -121,18 +136,27 @@ def add_tag(name, description):
 
 @tag.command('edit')
 @click.argument('name', type = TAG)
+@click.argument('implies', type = TAG, nargs = -1)
 @click.option('--description')
-def edit_tag(name, description):
-    from .model import Tag, Prefix, db
+def edit_tag(name, implies, description):
+    from .model import Tag, TagImplications, db
+    from peewee import IntegrityError
 
     try:
         tag = Tag.get(name = name.tag, prefix = name.prefix)
     except Tag.DoesNotExist:
         raise click.UsageError("Tag '%s' does not exist." % name)
 
-    if description:
-        tag.description = description
-        tag.save()
+    with db.atomic():
+        for i in fetch_tags(implies):
+            try:
+                TagImplications.create(tag = tag, implies_tag = i)
+            except IntegrityError:
+                print("Implication onto '%s' already existing. Skipping." % i)
+
+        if description:
+            tag.description = description
+            tag.save()
 
 @cli.group()
 def doc():
@@ -148,7 +172,6 @@ def add_doc(file, tags, create_tags, ignore_missing_tags):
     """Add a new document together with the given tags."""
 
     from .model import Document, Tag, db, DocumentTag
-    from .peewee_ext import sqlite_tuple_in
     import magic
 
     with db.atomic():
@@ -156,12 +179,7 @@ def add_doc(file, tags, create_tags, ignore_missing_tags):
             if create_tags:
                 tags = [create_tag(tag)[0] for tag in tags]
             else:
-                fetched_tags = Tag.select().where(sqlite_tuple_in((Tag.prefix, Tag.name), tags)).execute()
-                if len(fetched_tags) < len(tags):
-                    print("Some tags were not present:", ', '.join(set(map(str, tags)).difference(map(str, fetched_tags))))
-                    if not ignore_missing_tags:
-                        raise click.ClickException("Not all tags present")
-                tags = fetched_tags
+                tags = fetch_Tags(tags, ignore_missing_tags)
 
         mimetype = magic.from_file(file.name, mime=True)