diff --git a/src/main.py b/src/main.py index 72c5aa7..85471a0 100644 --- a/src/main.py +++ b/src/main.py @@ -10,6 +10,7 @@ import time def main(): ap = argparse.ArgumentParser() ap.add_argument("--stdin", default=False, action="store_true") + ap.add_argument("--stdout", default=False, action="store_true") ap.add_argument("--interval", default=250, type=int) ap.add_argument("--max-keys-down", default=2, type=int) ap.add_argument("--startup", default=3, type=int) @@ -24,6 +25,8 @@ def main(): w_translation = json.loads(args.translation) if args.translation else None w = writer.Writer(writer.PyAutoGUIWriter(w_translation)) + if args.stdout: + w = writer.Writer(writer.StdoutWriter()) if args.controller_setup: stdin = reader.StdinReader() diff --git a/src/reader.py b/src/reader.py index fd6ce3c..77cc5d1 100644 --- a/src/reader.py +++ b/src/reader.py @@ -69,16 +69,43 @@ class StdinReader: class RandomReader: def __init__(self, keys={"a":{"weight":1}, "b":{"weight":1}}): self.keys = keys - self.weighted_keys = [] - for k in self.keys: - weight = int(self.keys[k]["weight"] * 1000) - self.weighted_keys.extend([k] * weight) - log.info("weighted keys", self.weighted_keys) + self.pool = RandomPool(keys) def read(self): - return self.weighted_keys[ - random.randint(0, len(self.weighted_keys)-1) - ] + return self.pool.pop() + +class RandomPool: + def __init__(self, values): + self.values = {k:1000*v["weight"] for k,v in values.items()} + self.total = sum(self.values.values()) + self.consumed = set() + + def reset(self): + self.consumed = set() + + def pop(self): + k = self._pop() + for i in range(0, 3-1): + if k in self.consumed: + k = self._pop() + else: + break + self.consumed.add(k) + if self.should_reset(): + self.reset() + return k + + def _pop(self): + idx = random.randint(0, self.total-1) + idx_offset = 0 + for k in sorted(self.values.keys()): + if self.values[k] > idx+idx_offset: + return k + idx_offset -= self.values[k] + raise Exception(":(") + + def should_reset(self): + return len(self.consumed) >= len(self.values)/2 class FileReader: def __init__(self, path):