diff --git a/graph.py b/graph.py index 1568d09..6b9c974 100755 --- a/graph.py +++ b/graph.py @@ -49,10 +49,14 @@ if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--data', type=str, default=recipe_root) parser.add_argument('--target', type=str, nargs="*") + parser.add_argument('--prefer', type=str, nargs="*") args = parser.parse_args() recipes = satisfactory.Cookbook(args.data) + if args.prefer: + for alternative in args.prefer: + recipes.prefer(alternative) if args.target: graph(recipes, args.target) diff --git a/satisfactory/__init__.py b/satisfactory/__init__.py index 90e6003..2bd4a13 100644 --- a/satisfactory/__init__.py +++ b/satisfactory/__init__.py @@ -4,6 +4,8 @@ import logging from typing import Dict, Generator, Iterable, NamedTuple, Set, List, Optional +from collections import defaultdict + class Recipe(NamedTuple): input: Dict[str, int] @@ -33,6 +35,10 @@ class Item(NamedTuple): class Cookbook(object): __recipes: Dict[str, Item] + # Maps the name of an alternative recipe to the name of the component it + # applies to. + __alternatives: Dict[str, str] = defaultdict(lambda: []) + def __init__(self, root: str): self.__recipes = dict() @@ -51,6 +57,11 @@ class Cookbook(object): except Exception as ex: logging.error(f"Could not load {name}: {ex}") + for recipe in self.__recipes[name].recipes: + if recipe.requires: + self.__alternatives[recipe.requires].append(name) + + def __getitem__(self, item: str) -> Dict[str, Dict]: return self.__recipes[item] @@ -92,3 +103,14 @@ class Cookbook(object): :return: The names of all items that are resources. """ return (i for i in self.all() if self.is_resource(i)) + + def prefer(self, alternative: str): + def compare(x): + if not x.requires: + return 1 + if not x.requires == alternative: + return 1 + return 0 + + for recipe_name in self.__alternatives[alternative]: + self.__recipes[recipe_name].recipes.sort(key=compare)