plan: split the main function into component functions

plan.py

@@ -68,32 +68,21 @@ def required_rates(
 conveyor_rates = [0, 60, 120, 270, 480, 780, 900]
 
 
-def main():
+def required_machines(recipes: satisfactory.Cookbook, rates: Dict[str, Fraction]) -> Dict[str, int]:
-    recipes = satisfactory.Cookbook('data/recipes')
+    """
+    Calculate the number of machines required to build each item at the
+    requested rates. This will, by necessity, round up rates where they are
+    below the output rate of the relevant machine.
 
-    # Create a list of items we want to make
+    :param recipes: The cookbook object we're working from
-    # targets = [ 'supercomputer' ]
+    :param rates: A mapping from item names to requested output rates.
-    targets = recipes.components()
+    :return: A mapping from machine names to counts
-
+    """
-    # Create an initial name:rate request for all of the target items, then
-    # create a plan for their creation.
-    remain = [{n: basic_rate(recipes[n]['recipes'][0]) for n in targets}]
-    required_items = required_rates(recipes, remain)
-
-    # Note if any particular item is (in aggregate) going to exceed the
-    # highest conveyor belt capacity.
-    for name, rate in required_items.items():
-        print(name, rate, float(rate * 60))
-        if rate * 60 > conveyor_rates[-1]:
-            print("Rate exceeds max conveyor")
-
-    # Calculate the number of machines required to build each item at the
-    # calculated rates.
     def numberdict():
         return collections.defaultdict(int)
     required_machines = collections.defaultdict(numberdict)
 
-    for name, requested_rate in required_items.items():
+    for name, requested_rate in rates.items():
         if recipes.is_resource(name):
             continue
 
@@ -107,17 +96,54 @@ def main():
         machine = descriptor['machine']
         required_machines[machine][name] += requested_rate / normal_rate
 
-    # Calculate the power requirements for all the machines
+    return required_machines
-    required_power = 0
 
-    for machine, buckets in required_machines.items():
+
+def required_power(recipes: satisfactory.Cookbook, machines: Dict[str, int]) -> int:
+    """
+    Calculate the cumulative power requirements for a mapping of machine
+    names to counts
+
+    :param recipes: The cookbook we're working from
+    :param machines: A mapping from machine names to counts
+    :return: The total required power in MW
+    """
+    # Calculate the power requirements for all the machines
+    total = 0
+
+    for machine, buckets in machines.items():
         for result, rate in buckets.items():
             count = int(math.ceil(rate))
-            required_power += count * recipes[machine]['power_usage']
+            total += count * recipes[machine]['power_usage']
             print(machine, result, math.ceil(rate))
 
-    print(required_power)
+    return total
-    print(math.ceil(required_power / 150), "fuel generators")
+
+
+def main():
+    recipes = satisfactory.Cookbook('data/recipes')
+
+    # Create a list of items we want to make
+    # names = [ 'supercomputer' ]
+    names = recipes.components()
+
+    # Create an initial name:rate request for all of the target items, then
+    # create a plan for their creation.
+    remain = [{n: basic_rate(recipes[n]['recipes'][0]) for n in names}]
+    rates = required_rates(recipes, remain)
+
+    # Note if any particular item is (in aggregate) going to exceed the
+    # highest conveyor belt capacity.
+    for name, rate in rates.items():
+        print(name, rate, float(rate * 60))
+        if rate * 60 > conveyor_rates[-1]:
+            print("Rate exceeds max conveyor")
+
+    machines = required_machines(recipes, rates)
+    power = required_power(recipes, machines)
+
+    print(power, "MW")
+    print(math.ceil(power / 150), "fuel generators")
 
 
 if __name__ == '__main__':