Split util and tmsu functions

1 files changed, 71 insertions, 0 deletions

diff --git a/tmsu.py b/tmsu.py
new file mode 100644
index 0000000..4bd8946
--- /dev/null
+++ b/tmsu.py
@@ -0,0 +1,71 @@
+import subprocess
+import re
+import logging
+import os
+
+'''
+Initializes TMSU in a given directory.
+Does nothing, if it is already initialized.
+
+Parameters:
+base: Directory to initialize.
+'''
+def tmsu_init(base):
+    logger = logging.getLogger(__name__)
+    if not os.path.exists(os.path.join(base, ".tmsu")):
+        logger.info("TMSU database does not exist, creating ...")
+        proc = subprocess.Popen(["tmsu", "init"], cwd=base)
+        proc.wait()
+        logger.debug("TMSU returncode: {}".format(proc.returncode))
+        if proc.returncode != 0:
+            logger.error("Could not initialize TMSU database.")
+            return False
+    return True
+
+'''
+Reads the tags for the specified file.
+
+Parameters:
+base: Base directory for the database.
+file: File to get the tags for.
+'''
+def tmsu_tags(base, file):
+    logger = logging.getLogger(__name__)
+    logger.debug("Getting existing tags for file {}".format(file))
+    tags = set()
+    proc = subprocess.Popen(["tmsu", "tags", os.path.relpath(file, base)], cwd=base, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    proc.wait()
+    logger.debug("TMSU returncode: {}".format(proc.returncode))
+    if proc.returncode == 0:
+        ret = proc.stdout.read().decode()
+        logger.debug("Raw TMSU tags: {}".format(ret))
+        tags.update(re.split("\s", ret.split(":")[1])[1:-1])
+    else:
+        logger.error("Could not get tags for file {}".format(file))
+    return tags
+
+'''
+Sets tags for the specified file.
+
+Parameters:
+base: Base directory for the database.
+file: File to set the tags for.
+tags: Tags to set.
+untag: If True, it will remove all existing tags first. If False, it will just append new tags.
+'''
+def tmsu_tag(base, file, tags, untag=True):
+    logger = logging.getLogger(__name__)
+    if untag:
+        logger.debug("Untagging file")
+        proc = subprocess.Popen(["tmsu", "untag", "--all", os.path.relpath(file, base)], cwd=base, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        proc.wait()
+        if proc.returncode != 0:
+            logger.error("Could not untag file {}".format(file))
+    if tags:
+        logger.debug("Writing tags {}".format(tags))
+        proc = subprocess.Popen(["tmsu", "tag", os.path.relpath(file, base)] + list(tags), cwd=base, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        proc.wait()
+        if proc.returncode != 0:
+            logger.error("Could not write tags to file {}".format(file))
+    else:
+        logger.info("Tags are empty, ignoring")
\ No newline at end of file