TMSU: Classified and added option to override default command.

1 files changed, 65 insertions, 61 deletions

diff --git a/tmsu.py b/tmsu.py
index 4bd8946..aaa2947 100644
--- a/tmsu.py
+++ b/tmsu.py
@@ -3,69 +3,73 @@ import re
 import logging
 import os
 
-'''
-Initializes TMSU in a given directory.
-Does nothing, if it is already initialized.
+class TMSU(object):
 
-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
+    '''
+    Initializes TMSU in a given directory.
+    Does nothing, if it is already initialized.
 
-'''
-Reads the tags for the specified file.
+    Parameters:
+    base: Directory to initialize.
+    '''
+    def __init__(self, base, command):
+        self.__base = base
+        self.__command = command
+        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([self.__command, "init"], cwd=self.__base)
+            proc.wait()
+            logger.debug("TMSU returncode: {}".format(proc.returncode))
+            if proc.returncode != 0:
+                logger.error("Could not initialize TMSU database.")
+                self.status = False
+        self.status = True
 
-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
+    '''
+    Reads the tags for the specified file.
 
-'''
-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)
+    Parameters:
+    base: Base directory for the database.
+    file: File to get the tags for.
+    '''
+    def tags(self, file):
+        logger = logging.getLogger(__name__)
+        logger.debug("Getting existing tags for file {}".format(file))
+        tags = set()
+        proc = subprocess.Popen([self.__command, "tags", os.path.relpath(file, self.__base)], cwd=self.__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
+        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 tag(self, file, tags, untag=True):
+        logger = logging.getLogger(__name__)
+        if untag:
+            logger.debug("Untagging file")
+            proc = subprocess.Popen([self.__command, "untag", "--all", os.path.relpath(file, self.__base)], cwd=self.__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([self.__command, "tag", os.path.relpath(file, self.__base)] + list(tags), cwd=self.__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