Add labs to logic

base/data/classes.py

@@ -111,6 +111,7 @@ class Machine(Base):
     surface_conditions: list[SurfaceCondition] = dataclasses.field(default_factory=list)
     crafting_categories: set[str] = dataclasses.field(default_factory=set)
     mining_categories: set[str] = dataclasses.field(default_factory=set)
+    allowed_science_packs: set[str] = dataclasses.field(default_factory=set)
     is_offshore_pump: bool = False
     is_asteroid_collector: bool = False
 

base/data/lookup.py

@@ -6,6 +6,7 @@ def machines_by(
     can_be_placed_on: Surface|None = None,
     crafting_category: str|None = None,
     mining_category: str|None = None,
+    allowed_science_packs: set[str]|None = None,
     is_offshore_pump: bool = None,
     is_asteroid_collector: bool = None,
 ) -> list[Machine]:
@@ -16,6 +17,8 @@ def machines_by(
         filtered_machines = filter(lambda machine: crafting_category in machine.crafting_categories, filtered_machines)
     if mining_category is not None:
         filtered_machines = filter(lambda machine: mining_category in machine.mining_categories, filtered_machines)
+    if allowed_science_packs is not None:
+        filtered_machines = filter(lambda machine: machine.allowed_science_packs.issuperset(allowed_science_packs), filtered_machines)
     if is_offshore_pump is not None:
         filtered_machines = filter(lambda machine: machine.is_offshore_pump, filtered_machines)
     if is_asteroid_collector is not None:

base/data/raw_base.py

@@ -181,6 +181,13 @@ def _recursive_asteroid_to_chunks(asteroid_name: str):
         set(prototype['crafting_categories']),
     ))
 
+for prototype in get_prototypes('lab'):
+    machines.add(Machine(
+        prototype['name'],
+        parse_surface_condition_array(prototype.get('surface_conditions', [])),
+        allowed_science_packs=set(prototype['inputs']),
+    ))
+
 for prototype in get_prototypes('mining-drill'):
     machines.add(Machine(
         prototype['name'],

base/world/locations/factory.py

@@ -1,7 +1,7 @@
 from random import Random
 
 from ...config import items_required_for_automation, items_required_for_research
-from ...data.lookup import recipes_by
+from ...data.lookup import machines_by, recipes_by
 from ...data.raw import science_packs, surfaces
 from ...data.utils import craftable_items_at_start
 from ..options import FactorioOptions
@@ -90,6 +90,12 @@ def get_locations(options: FactorioOptions, random: Random, location_count: int)
             All([
                 HasProduction(item_name, surface, i >= early_science_location_count)
                 for item_name in science_location.ingredients
+            ]) & Any([
+                HasProduction(lab.name, surface)
+                for lab in machines_by(
+                    can_be_placed_on=surface,
+                    allowed_science_packs=science_location.ingredients.keys(),
+                )
             ])
             for surface in surfaces
         ])