From 111fc49ea8b8d3be56be26362beb2273a2291c1b Mon Sep 17 00:00:00 2001 From: Yorick van Pelt Date: Mon, 8 Jan 2024 12:06:10 +0100 Subject: [PATCH] have some types --- magisch_corvee_script.py | 64 ++++++++++++++++++++++++---------------- 1 file changed, 38 insertions(+), 26 deletions(-) diff --git a/magisch_corvee_script.py b/magisch_corvee_script.py index c29a1f5..05db27f 100755 --- a/magisch_corvee_script.py +++ b/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'])] -ALL_DAYS = set(range(config['days'])) + [sum(task.workload * task.req[d] for task in tasks.values()) for d in 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 isdict(x): - return type(x) == dict +def read_people(conf_ppl) -> dict[str, Person]: 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']: - task_lut[lookup] = t + for lookup in taskconf.lookup: + 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): - for ti,(task,conf) in enumerate(tasks.items()): - if conf['personen'] == 'iedereen': +def set_capabilities(tasks: dict[str, TaskConfig], people: dict[str, Person]): + for (task,conf) in tasks.items(): + 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: - for day, pers in enumerate(conf['hardcode']): + if conf.hardcode is not None: + 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] - workload = {tsk_idx[t]: tasks[t]['workload'] for t in tasks} - return [loves, hates, capab, hardcode, max_loads, req, workload, costs] + req[di][ti] = t.req[di] + workload = {tsk_idx[t]: tasks[t].workload for t in tasks} + 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)