krebsdance: add flag to generate directed graph

krebs/5pkgs/simple/krebsdance/default.nix

@@ -2,113 +2,152 @@
 writers.writePython3Bin "krebsdance" {} ''
   import argparse
   import random
+  import itertools
 
   claws = [
       dict(
-          up='(\\/)',
+          up="(\\/)",
-          down='(/\\)',
+          down="(/\\)",
-          left='(\\\\)',
+          left="(\\\\)",
-          right='(//)',
+          right="(//)",
       ),
       dict(
-          up='(V)',
+          up="(V)",
-          down='(A)',
+          down="(A)",
-          left='>)=',
+          left=">)=",
-          right='=(<',
+          right="=(<",
       ),
       dict(
-          up='(U)',
+          up="(U)",
-          down='(n)',
+          down="(n)",
-          left=')==',
+          left=")==",
-          right='==(',
+          right="==(",
       ),
   ]
 
   eyes = [
-      '°',
+      "°",
-      '*',
+      "*",
-      '^',
+      "^",
-      'ö',
+      "ö",
-      'o',
+      "o",
-      'O',
+      "O",
-      'X',
+      "X",
-      'x',
+      "x",
-      'U',
+      "U",
-      'u',
+      "u",
   ]
 
   bodies = [
       dict(
-          left='(',
+          left="(",
-          right=')',
+          right=")",
       ),
       dict(
-          left='{',
+          left="{",
-          right='}',
+          right="}",
       ),
       dict(
-          left='[',
+          left="[",
-          right=']',
+          right="]",
       ),
       dict(
-          left='<',
+          left="<",
-          right='>',
+          right=">",
       ),
       dict(
-          left='|',
+          left="|",
-          right='|',
+          right="|",
       ),
   ]
 
   mouths = [
-      ',,,,',
+      ",,,,",
-      ',mm,',
+      ",mm,",
-      '_mm_',
+      "_mm_",
-      '-mm-',
+      "-mm-",
-      ';;;;',
+      ";;;;",
-      ';mm;',
+      ";mm;",
-      ':mm:',
+      ":mm:",
-      '::::',
+      "::::",
-      ':ww:',
+      ":ww:",
-      ':<>:',
+      ":<>:",
   ]
 
 
+  def all_krebses():
+      for mouth, body, eye, claw in itertools.product(mouths, bodies, eyes, claws):
+          yield f'{claw["up"]} {body["left"]}{eye}{mouth}{eye}{body["right"]} {claw["up"]}'
+
+
+  def krebs_graph() -> str:
+      return "\n".join(
+          ["digraph {"]
+          + [f'"{krebs}"->"{generate(seed=krebs)}"' for krebs in all_krebses()]
+          + ["}"]
+      )
+
+
+  def generate(*, seed: str, dancing: bool = False) -> str:
+      if seed:
+          random.seed(seed)
+      clawstyle = random.choice(claws)
+      body = random.choice(bodies)
+      eye = random.choice(eyes)
+      mouth = random.choice(mouths)
+      if dancing:
+          return "\n".join(
+              [
+                  f'{clawstyle["down"]} {body["left"]}{eye}{mouth}{eye}{body["right"]}{clawstyle["up"]}',
+                  f'{clawstyle["left"]}{body["left"]}{eye}{mouth}{eye}{body["right"]} {clawstyle["right"]}',
+                  f'{clawstyle["right"]} {body["left"]}{eye}{mouth}{eye}{body["right"]} {clawstyle["left"]}',
+                  f'{clawstyle["down"]}{body["left"]}{eye}{mouth}{eye}{body["right"]}{clawstyle["down"]}',
+              ]
+          )
+      else:
+          return f'{clawstyle["up"]} {body["left"]}{eye}{mouth}{eye}{body["right"]} {clawstyle["up"]}'
+
+
+  def fixpoints():
+      for krebs in all_krebses():
+          if generate(seed=krebs) == krebs:
+              yield krebs
+
+
   def main():
       parser = argparse.ArgumentParser()
 
       parser.add_argument(
-          'seed',
+          "seed",
-          nargs='?',
+          nargs="?",
-          help='random seed to use for generating the krebs variant',
+          help="random seed to use for generating the krebs variant",
       )
 
       parser.add_argument(
-          '--dance', '-d',
+          "--dance",
-          dest='dance',
+          "-d",
-          help='if the krebs should dance',
+          dest="dance",
+          help="if the krebs should dance",
           default=False,
-          action='store_true',
+          action="store_true",
+      )
+
+      parser.add_argument(
+          "--mode",
+          "-m",
+          dest="mode",
+          choices=["graphviz", "plain"],
+          default="plain",
       )
 
       args = parser.parse_args()
 
-      if args.seed:
+      if args.mode == "plain":
-          random.seed(args.seed)
+          print(generate(seed=args.seed, dancing=args.dance))
-
+      elif args.mode == "graphviz":
-      clawstyle = random.choice(claws)
+          print(krebs_graph())
-      body = random.choice(bodies)
-      eye = random.choice(eyes)
-      mouth = random.choice(mouths)
-      if args.dance:
-          print(f'{clawstyle["down"]} {body["left"]}{eye}{mouth}{eye}{body["right"]}{clawstyle["up"]}')  # noqa
-          print(f' {clawstyle["left"]}{body["left"]}{eye}{mouth}{eye}{body["right"]} {clawstyle["right"]}')  # noqa
-          print(f'{clawstyle["right"]} {body["left"]}{eye}{mouth}{eye}{body["right"]} {clawstyle["left"]}')  # noqa
-          print(f' {clawstyle["down"]}{body["left"]}{eye}{mouth}{eye}{body["right"]}{clawstyle["down"]}')  # noqa
-      else:
-          print(f'{clawstyle["up"]} {body["left"]}{eye}{mouth}{eye}{body["right"]} {clawstyle["up"]}')  # noqa
 
 
-  if __name__ == '__main__':
+  if __name__ == "__main__":
       main()
 ''