Handle borg output in separate class

src/borgoutputhandler.py

@@ -0,0 +1,24 @@
+from datetime import datetime
+import borg
+import json
+
+
+class BorgOutputHandler(object):
+    def __init__(self, borg_output: str):
+        self.borg_output = borg_output
+        self.borg_json = None
+        self.error = False
+        try:
+            self.borg_json = json.loads(borg_output)
+        except json.JSONDecodeError:
+            self.error = True
+
+    def get_borg_info(self):
+        repo = borg.Repo.from_json(self.borg_json['repository'])
+        archive = borg.Archive.from_json(self.borg_json['archive'])
+        stats = borg.Stats.from_json(self.borg_json['archive']['stats'])
+
+        return repo, archive, stats
+
+    def get_borg_error(self):
+        return borg.Error(self.borg_output, datetime.now())

src/database/borgdatabase.py

@@ -1,8 +1,5 @@
 from .connection import RepoConn, ArchiveConn, StatsConn, ErrorConn
-from datetime import datetime
 from pathlib import Path
-from src import borg
-import json
 
 
 class BorgDatabase(object):
@@ -20,24 +17,10 @@ class BorgDatabase(object):
         self.error_conn = ErrorConn(db_path,
                                     table_name=self.error_name)
 
-    def process_borg_output(self, borg_output: str):
+    def insert_record(self, repo, archive, stats):
-        borg_json = None
-        try:
-            borg_json = json.loads(borg_output)
-        except json.JSONDecodeError:
-            self.handle_borg_error(borg_output)
-            return
-        self.process_borg_json(borg_json)
-
-    def process_borg_json(self, borg_json: dict):
-        repo = borg.Repo.from_json(borg_json['repository'])
-        archive = borg.Archive.from_json(borg_json['archive'])
-        stats = borg.Stats.from_json(borg_json['archive']['stats'])
-
         repo_id = self.repo_conn.insert(repo)
         archive_id = self.archive_conn.insert(archive, repo_id)
         self.stats_conn.insert(stats, repo_id, archive_id)
 
-    def handle_borg_error(self, borg_error: str):
+    def insert_error(self, borg_error):
-        error = borg.Error(borg_error, datetime.now())
+        self.error_conn.insert(borg_error)
-        self.error_conn.insert(error)

src/main.py

@@ -2,16 +2,32 @@ from database import BorgDatabase
 from sys import stdin
 from os.path import realpath
 from pathlib import Path
+import argparse
+from borgoutputhandler import BorgOutputHandler
 
 
-def main(borg_output: str, path: Path):
+def main(args, path: Path):
-    db_path = path / 'borg.sqlite'
+    if args.graph is not None:
-    db = BorgDatabase(db_path)
+        pass
+    else:
+        borg_output = " ".join(stdin.readlines())
+        bo = BorgOutputHandler(borg_output)
+        db = BorgDatabase(path / 'borg.sqlite')
 
-    db.process_borg_output(borg_output)
+        if bo.error:
+            db.insert_error(bo.get_borg_error())
+        else:
+            db.insert_record(*bo.get_borg_info())
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-g", "--graph", help="Produce graphs at specified location",
+                        type=str)
+    return parser.parse_args()
 
 
 if __name__ == "__main__":
+    args = get_args()
     path = Path(realpath(__file__)).parent.parent
-    borg_output = "\n".join(stdin.readlines())
+    main(args, path)
-    main(borg_output, path)