satisfactory/graph.py

#!/usr/bin/env python3

"""
This script takes a target item and produces GraphViz formatted output that
describes the gross production dependencies required to produce the item.
"""

import satisfactory

import os

from typing import Iterable


def graph(cookbook: satisfactory.Cookbook, targets: Iterable[str]):
    """
    Output the directed graph representing the production chain needed to
    produce all the provided targets.

    :param cookbook:
    :param targets:
    :return:
    """
    print("digraph G {")

    seen = set()
    remain = set(targets)

    while remain:
        output = remain.pop()

        for need in cookbook.recipes(output)[0].input:
            print(f"{need} -> {output}")
            if need not in seen:
                remain.add(need)

        seen.add(output)

    print("}")


if __name__ == '__main__':
    def main():
        import argparse

        root = os.path.dirname(__file__)
        recipe_root = os.path.join(root, 'data', 'recipes')

        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)
        else:
            graph(recipes, recipes.components())

    main()