have some types

magisch_corvee_script.py

@@ -6,14 +6,28 @@ import re
 import subprocess
 from collections import OrderedDict
 import glpm
+from typing import Any, Tuple, TypeVar
+from dataclasses import dataclass, field
 conf = yaml.safe_load(open('config.yaml', 'r'))
 config = conf['config']
 config['ignore'].append('')
-tasks = OrderedDict(conf['tasks'])
+
-index = lambda x: {v:k for k,v in enumerate(x)}
+@dataclass
+class TaskConfig:
+    personen: list[str]
+    workload: int
+    req: list[int]
+    hardcode: list[str] | None = None
+    lookup: list[str] = field(default_factory=list)
+
+tasks: dict[str, TaskConfig] = OrderedDict({k: TaskConfig(**t) for k, t in conf['tasks'].items()})
+X = TypeVar("X")
+def index(x: dict[X, Any]) -> dict[X, int]:
+    return {v: k for k, v in enumerate(x)}
+
 daily_workloads = \
-     [sum(task['workload'] * task['req'][d] for task in tasks.values()) for d in range(config['days'])]
+     [sum(task.workload * task.req[d] for task in tasks.values()) for d in range(config['days'])]
-ALL_DAYS = set(range(config['days']))
+ALL_DAYS: set[int] = set(range(config['days']))
 class Person(object):
     def __init__(self, conf={"dagen":ALL_DAYS}):
         self.can = set()
@@ -25,20 +39,18 @@ class Person(object):
         self.conf['dagen'] = set(conf['dagen'])
     def vrolijkheid(self):
         res = config['days'] - len(self.does)
-        for (d,t) in self.does:
+        for (_,t) in self.does:
             if t in self.loves:
                 res += config['weights']['likes']
             if t in self.hates:
                 res -= config['weights']['hates']
         return res
     def workload(self, tasks):
-        return sum(tasks[t]['workload'] for (d,t) in self.does)
+        return sum(tasks[t]['workload'] for (_,t) in self.does)
     def cost(self, num_people):
         return round(sum((daily_workloads[d] for d in self.conf['dagen'])) / num_people)
         # probabilistic round: int(math.floor(x + random.random()))
-def read_people(conf_ppl):
+def read_people(conf_ppl) -> dict[str, Person]:
-    def isdict(x):
-        return type(x) == dict
     people = OrderedDict()
     for x in conf_ppl:
         val = {"dagen": ALL_DAYS}
@@ -47,12 +59,11 @@ def read_people(conf_ppl):
         people[x.lower()] = Person(val)
     return people
 # deal with loves/hates
-def make_task_lut(tasks):
+def make_task_lut(tasks: dict[str, TaskConfig]):
     task_lut = {}
     for t, taskconf in tasks.items():
-        if 'lookup' in taskconf:
+        for lookup in taskconf.lookup:
-            for lookup in taskconf['lookup']:
+            task_lut[lookup] = t
-                task_lut[lookup] = t
     task_re = re.compile(config['task_re'])
     def lookup_tasks(tasks):
         return (task_lut[x] for x in task_re.split(tasks) if not x in config['ignore'])
@@ -69,23 +80,23 @@ def read_prefs(pref_file, tasks, people):
         if not p.has_prefs:
             print("warning: no preferences for", name, file=sys.stderr)
 # deal with capability and hardcode
-def set_capabilities(tasks, people):
+def set_capabilities(tasks: dict[str, TaskConfig], people: dict[str, Person]):
-    for ti,(task,conf) in enumerate(tasks.items()):
+    for (task,conf) in tasks.items():
-        if conf['personen'] == 'iedereen':
+        if conf.personen == 'iedereen':
             for p in people.values():
                 p.can.add(task)
-        elif conf['personen'] == 'liefhebbers':
+        elif conf.personen == 'liefhebbers':
             for p in people.values():
                 if task in p.loves:
                     p.can.add(task)
         else:
-            for p in conf['personen']:
+            for p in conf.personen:
                 people[p.lower()].can.add(task)
-        if 'hardcode' in conf:
+        if conf.hardcode is not None:
-            for day, pers in enumerate(conf['hardcode']):
+            for day, pers in enumerate(conf.hardcode):
                 people[pers.lower()].does.add((day, task))
 # format as matrices
-def matrices(people, tasks):
+def matrices(people: dict[str, Person], tasks: dict[str, TaskConfig]) -> Tuple[list[list[int]], list[list[int]], list[list[int]], dict[str, int], dict[Tuple[int, int], int], list[list[int]], dict[int, int], dict[int, int]]:
     mat = lambda a,b: [[0 for j in b] for i in a]
     loves = mat(people, tasks)
     hates = mat(people, tasks)
@@ -94,7 +105,7 @@ def matrices(people, tasks):
     hardcode = {}
     max_loads = {}
     costs = {}
-    for i,(person, p) in enumerate(people.items()):
+    for i,p in enumerate(people.values()):
         for t in p.loves: loves[i][tsk_idx[t]] = 1
         for t in p.hates: hates[i][tsk_idx[t]] = 1
         for t in p.can: capab[i][tsk_idx[t]] = 1
@@ -109,9 +120,9 @@ def matrices(people, tasks):
     req = mat(range(config['days']), tasks)
     for di in range(config['days']):
         for ti,t in enumerate(tasks.values()):
-            req[di][ti] = t['req'][di]
+            req[di][ti] = t.req[di]
-    workload = {tsk_idx[t]: tasks[t]['workload'] for t in tasks}
+    workload = {tsk_idx[t]: tasks[t].workload for t in tasks}
-    return [loves, hates, capab, hardcode, max_loads, req, workload, costs]
+    return (loves, hates, capab, hardcode, max_loads, req, workload, costs)
 def read_assignment(file, people, tasks):
     def between_the_lines(f, a=">>>>\n", b="<<<<\n"):
         for l in f:
@@ -131,7 +142,7 @@ def write_data(people, tasks, file=sys.stdout):
     print(glpm.matrix("L", loves), file=file)
     print(glpm.matrix("H", hates), file=file)
     print(glpm.matrix("C", capab, 1), file=file)
-    print(glpm.matrix("R", req, None), file=file)
+    print(glpm.matrix("R", req, 0), file=file)
     print(glpm.dict("Q", hardcode), file=file)
     print(glpm.dict("Wl", workload), file=file)
     print(glpm.dict("max_load", max_loads), file=file)
@@ -152,6 +163,7 @@ def write_tasks(people, tasks, file=sys.stdout):
         days_filled = days_fmt.format(*map(q, days))
         print("| {} ||{} {} || {}".format(name, days_filled, p.vrolijkheid(), p.workload(tasks)), file=file)
         print("|-")
+
 people = read_people(conf['people'])
 with open('prefs_table', 'r') as pref_file:
     read_prefs(pref_file, tasks, people)