From 7c0ccaf0b6aceb95dda97924b83f9e91d3955b23 Mon Sep 17 00:00:00 2001
From: Yorick van Pelt <yorick@yorickvanpelt.nl>
Date: Mon, 15 Apr 2024 10:22:47 +0200
Subject: [PATCH] Use tabulate for output format, add day names and --simple

---
 magisch_corvee_script.py | 34 ++++++++++++++++++++++++----------
 requirements.txt         |  3 +++
 2 files changed, 27 insertions(+), 10 deletions(-)
 create mode 100644 requirements.txt

diff --git a/magisch_corvee_script.py b/magisch_corvee_script.py
index 3b82a71..22fe847 100755
--- a/magisch_corvee_script.py
+++ b/magisch_corvee_script.py
@@ -9,13 +9,16 @@ from collections import OrderedDict, defaultdict
 from pyscipopt import Model, quicksum
 from typing import Any, Tuple, TypeVar
 from dataclasses import dataclass, field
+from tabulate import tabulate
 conf = yaml.safe_load(open('config.yaml', 'r'))
 
 DEFAULT_CONFIG = {
-    "max_load_person": 6
+    "max_load_person": 6,
+    "day_names": [f"dag {str(i)}" for i in range(conf['config']['days'])],
 }
 config = DEFAULT_CONFIG | conf['config']
 config['ignore'].append('')
+assert(len(config['day_names']) == config['days'])
 
 QUADRATIC = False
 
@@ -105,17 +108,26 @@ def set_capabilities(tasks: dict[str, TaskConfig], people: dict[str, Person]):
                 if pers:
                     people[pers.lower()].does.add((day, task))
 
-def write_tasks(people, tasks, file=sys.stdout):
-    for name, p in people.items():
-        days = [[] for i in range(config['days'])]
+def write_tasks(people: dict[str, Person], tasks: dict[str, TaskConfig], file=sys.stdout):
+    headers = ["wie"] + config['day_names']
+    if not args.simple:
+        headers += ["workload", "vrolijkheid"]
+    tabl = []
+    for p in people.values():
+        days = [[] for _ in range(config['days'])]
         for (d,t) in p.does:
             days[d].append((t, t in p.loves, t in p.hates))
-        def q(w):
-            return ",".join([t + (" <3" if love else "") + (" :(" if hate else "") for (t,love,hate) in w])
-        days_fmt = " {} ||" * len(days)
-        days_filled = days_fmt.format(*map(q, days))
-        print("| {} ||{} {} || {}".format(name, days_filled, p.vrolijkheid(), p.workload(tasks)), file=file)
-        print("|-")
+        if not args.simple:
+            def q(w):
+                return ",".join([tasks[t].name + (" :)" if love else "") + (" :(" if hate else "") for (t,love,hate) in w])
+        else:
+            def q(w):
+                return ",".join([tasks[t].name for (t,_,_) in w])
+        row = [p.name, *map(q, days)]
+        if not args.simple:
+            row += [p.workload(tasks), p.vrolijkheid()]
+        tabl.append(row)
+    print(tabulate(tabl, headers=headers, tablefmt=args.output_format), file=file)
 
 def scipsol(people: dict[str, Person], tasks: dict[str, TaskConfig]):
     max_loads: dict[Tuple[str, int], int] = {}
@@ -201,7 +213,9 @@ def scipsol(people: dict[str, Person], tasks: dict[str, TaskConfig]):
 
 parser = argparse.ArgumentParser()
 parser.add_argument("-q", "--quadratic", action="store_true")
+parser.add_argument("--simple", action="store_true", help="hide workload and happiness")
 parser.add_argument("--max_total_error", type=int, default=None)
+parser.add_argument("--output-format", default="mediawiki", help="`tabulate` output format")
 
 args = parser.parse_args()
 QUADRATIC = args.quadratic
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..6a3f0c8
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,3 @@
+tabulate~=0.9.0
+PySCIPOpt~=4.4.0
+PyYAML~=6.0.1