From 3ae3dfe9225ac7317725707f3dcb4c65fb93867b Mon Sep 17 00:00:00 2001
From: Emi Vasilek <emi.vasilek@gmail.com>
Date: Mon, 20 Nov 2023 10:27:21 +0100
Subject: [PATCH] add generate-units and generate-ingredients subcommands

---
 comfyrecipes/builder.py | 77 ++++++++++++++++++++++++++++++++++++++++-
 comfyrecipes/cli.py     | 25 +++++++++++--
 2 files changed, 99 insertions(+), 3 deletions(-)

diff --git a/comfyrecipes/builder.py b/comfyrecipes/builder.py
index c5f1f29..0db6d9e 100644
--- a/comfyrecipes/builder.py
+++ b/comfyrecipes/builder.py
@@ -1,4 +1,5 @@
-from typing import Any, List, Set
+import sys
+from typing import Any, Callable, Dict, List, Set, TypeVar
 import os
 import json
 
@@ -9,6 +10,7 @@ import jsonschema.exceptions
 
 import comfyrecipes.parsing as parsing
 
+T = TypeVar("T")
 
 class Builder:
     def __init__(self) -> None:
@@ -132,6 +134,79 @@ class Builder:
             )
         return 0
 
+    def generate_units(self, dir: str) -> None:
+        def collect_unitnames(rec: parsing.Recipe) -> List[str]:
+            results: List[str] = []
+            for ing in rec.ingredients:
+                results.append(ing.unit.name)
+            return results
+        unitnamelists = self.foreach_recipe(dir, collect_unitnames)
+        unitnamesset: Set[str] = set()
+        for unitnamelst in unitnamelists:
+            for unitname in unitnamelst:
+                unitnamesset.add(unitname)
+        unitnames = list(unitnamesset)
+        unitnames.sort()
+
+        units: List[Dict[str, str]] = []
+        for name in unitnames:
+            units.append({"name": name})
+
+        file = f"{dir}/units.yaml"
+        with open(file, "w") as f:
+            f.write(yaml.dump(units))
+        print("found units written to", file)
+
+    def generate_ingredients(self, dir: str) -> None:
+        def collect_ingnames(rec: parsing.Recipe) -> List[str]:
+            results: List[str] = []
+            for ing in rec.ingredients:
+                results.append(ing.name)
+            return results
+        ingredientnamelists = self.foreach_recipe(dir, collect_ingnames)
+        ingredientnamesset: Set[str] = set()
+        for ingredientnamelst in ingredientnamelists:
+            for ingredientname in ingredientnamelst:
+                ingredientnamesset.add(ingredientname)
+        ingredientnames = list(ingredientnamesset)
+        ingredientnames.sort()
+
+        ingredients: List[Dict[str, str]] = []
+        for name in ingredientnames:
+            ingredients.append({"name": name})
+
+        file = f"{dir}/ingredients.yaml"
+        with open(file, "w") as f:
+            f.write(yaml.dump(ingredients))
+        print("found ingredients written to", file)
+
+    def foreach_recipe(self, dir: str, func: Callable[[parsing.Recipe], T]) -> List[T]:
+        files = []
+        for _, _, filesx in os.walk(dir + "/recipes"):
+            files = filesx
+            files.sort()
+
+        def foreach_subrecipe(recipe: parsing.Recipe) -> List[T]:
+            results: List[T] = []
+            for rec in [recipe] + recipe.subrecipes:
+                results.append(func(rec))
+            for subrec in recipe.subrecipes:
+                for subsubrec in subrec.subrecipes:
+                    results += foreach_subrecipe(subsubrec)
+            return results
+
+        results: List[T] = []
+        for file in files:
+            if not file.endswith(".yaml"):
+                print(f"unknown extension of {file}")
+                continue
+            recipedct = self.load_file(dir + "/recipes/" + file)
+            recipe = parsing.Recipe.from_dict(self.ctx, recipedct)
+            if self.ctx.issues.check() != 0:
+                continue
+            results += foreach_subrecipe(recipe)
+        return results
+
     def finish(self, dir: str) -> int:
         files = set()
         for _, _, filesx in os.walk(f"{dir}/out/html"):
diff --git a/comfyrecipes/cli.py b/comfyrecipes/cli.py
index 0240f09..e7f9173 100644
--- a/comfyrecipes/cli.py
+++ b/comfyrecipes/cli.py
@@ -4,7 +4,7 @@ import socketserver
 import sys
 import os
 
-import comfyrecipes.builder as builder
+import comfyrecipes.builder
 
 
 def main() -> None:
@@ -19,11 +19,20 @@ def main() -> None:
     parser_serve.add_argument("--port", type=int, default=8000)
     parser_serve.add_argument("--address", type=str, default="127.0.0.1")
 
+    parser_generate_units = subparsers.add_parser("generate-units")
+    parser_generate_units.add_argument("directory", type=str)
+    parser_generate_units.add_argument("--force", dest="force", action="store_true")
+
+    parser_generate_ingredients = subparsers.add_parser("generate-ingredients")
+    parser_generate_ingredients.add_argument("directory", type=str)
+    parser_generate_ingredients.add_argument("--force", dest="force", action="store_true")
+
     args = parser.parse_args()
 
     ret = 0
+    builder = comfyrecipes.builder.Builder()
     if args.subcommand == "build":
-        ret = builder.Builder().build(args.directory)
+        ret = builder.build(args.directory)
     elif args.subcommand == "serve":
         os.chdir(f"{args.directory}/out/html")
         httpd = socketserver.TCPServer(
@@ -34,6 +43,18 @@ def main() -> None:
             httpd.serve_forever()
         except KeyboardInterrupt:
             pass
+    elif args.subcommand == "generate-units":
+        if not args.force and os.path.isfile(args.directory + "/units.yaml"):
+            print("units.yaml already exists, pass --force if you want to overwrite it", file=sys.stderr)
+            ret = 1
+        else:
+            builder.generate_units(args.directory)
+    elif args.subcommand == "generate-ingredients":
+        if not args.force and os.path.isfile(args.directory + "/ingredients.yaml"):
+            print("ingredients.yaml already exists, pass --force if you want to overwrite it", file=sys.stderr)
+            ret = 1
+        else:
+            builder.generate_ingredients(args.directory)
     else:
         # unhandled, but valid subcommand
         assert False