From 67fcd6b8234abd5bdaee6ba4c0ecc756bc27eedb Mon Sep 17 00:00:00 2001
From: lepideble <147614625+lepideble@users.noreply.github.com>
Date: Sun, 12 Jul 2026 16:14:41 +0200
Subject: [PATCH] Add labs to logic

---
 base/data/classes.py            | 1 +
 base/data/lookup.py             | 3 +++
 base/data/raw_base.py           | 7 +++++++
 base/world/locations/factory.py | 8 +++++++-
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/base/data/classes.py b/base/data/classes.py
index 42156bc..ef8b114 100644
--- a/base/data/classes.py
+++ b/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
 
diff --git a/base/data/lookup.py b/base/data/lookup.py
index 9eae6de..094bada 100644
--- a/base/data/lookup.py
+++ b/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:
diff --git a/base/data/raw_base.py b/base/data/raw_base.py
index e1b342c..8628ed6 100644
--- a/base/data/raw_base.py
+++ b/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'],
diff --git a/base/world/locations/factory.py b/base/world/locations/factory.py
index c594ca4..1618080 100644
--- a/base/world/locations/factory.py
+++ b/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
         ])