diff --git a/graph.py b/graph.py index dabb5d0..5595ec4 100755 --- a/graph.py +++ b/graph.py @@ -3,21 +3,35 @@ import json import os +from typing import Iterable -def graph_all(recipes: dict) -> None: - print("digraph G {") - for result, method in recipes.items(): - for i in method[0]["input"]: - print(f"{i} -> {result}") - print("}") +class Cookbook(object): + recipes: dict + + def __init__(self, root: str): + self.recipes = dict() + + for dirname, dirs, files in os.walk(root): + for f in files: + path = os.path.join(dirname, f) + name, _ = os.path.splitext(f) + with open(path, 'r') as src: + variations = json.load(src) + + self.recipes[name] = variations + + def __getitem__(self, item: str): + return self.recipes[item] + + def all(self): + return self.recipes.keys() -def graph_one(recipes: dict, target: str) -> None: +def graph(recipes: dict, targets: Iterable[str]): print("digraph G {") seen = set() - remain = set() - remain.add(target) + remain = set(targets) while remain: output = remain.pop() @@ -31,19 +45,12 @@ def graph_one(recipes: dict, target: str) -> None: print("}") -def load_recipes(recipe_root: str) -> dict: - recipes = dict() +def graph_all(recipes: dict) -> None: + graph(recipes, recipes.all()) - for dirname, dirs, files in os.walk(recipe_root): - for f in files: - path = os.path.join(dirname, f) - name, _ = os.path.splitext(f) - with open(path, 'r') as src: - variations = json.load(src) - recipes[name] = variations - - return recipes +def graph_one(recipes: dict, target: str) -> None: + graph(recipes, [target]) if __name__ == '__main__': @@ -59,7 +66,7 @@ if __name__ == '__main__': args = parser.parse_args() - recipes = load_recipes(args.data) + recipes = Cookbook(args.data) if args.target: graph_one(recipes, args.target)