[Draft] Literal Include

Quick draft w/ Cursor on include support.

Author: ax3l

src/pals/functions.py

@@ -3,15 +3,15 @@
 import os
 
 
-def inspect_file_extensions(filename: str):
+def inspect_file_extensions(filename: str, check_extension: bool = True):
     """Attempt to strip two levels of file extensions to determine the schema.
 
     filename examples: fodo.pals.yaml, fodo.pals.json, ...
     """
     file_noext, extension = os.path.splitext(filename)
     file_noext_noext, extension_inner = os.path.splitext(file_noext)
 
-    if extension_inner != ".pals":
+    if check_extension and extension_inner != ".pals":
         raise RuntimeError(
             f"inspect_file_extensions: No support for file {filename} with extension {extension}. "
             f"PALS files must end in .pals.json or .pals.yaml or similar."
@@ -25,11 +25,94 @@ def inspect_file_extensions(filename: str):
     }
 
 
-def load_file_to_dict(filename: str) -> dict:
+def process_includes(data, base_dir: str):
+    """Recursively process 'include' directives in the data structure."""
+    if isinstance(data, dict):
+        # Handle 'include' key in dictionary
+        if "include" in data:
+            include_file = data["include"]
+            # Check if include_file is a string (filename)
+            if isinstance(include_file, str):
+                filepath = os.path.join(base_dir, include_file)
+                # Load included file without strict extension check
+                included_data = load_file_to_dict(filepath, check_extension=False)
+
+                # Remove 'include' key
+                local_data = data.copy()
+                del local_data["include"]
+
+                # Recursively process local data
+                local_data = {
+                    k: process_includes(v, base_dir) for k, v in local_data.items()
+                }
+
+                # Merge logic
+                # If included data is a list of single-key dicts (PALS special case), try to merge as dict
+                if isinstance(included_data, list):
+                    try:
+                        merged_included = {}
+                        all_dicts = True
+                        for item in included_data:
+                            if isinstance(item, dict) and len(item) == 1:
+                                merged_included.update(item)
+                            else:
+                                all_dicts = False
+                                break
+                        if all_dicts:
+                            included_data = merged_included
+                    except Exception:
+                        pass
+
+                if isinstance(included_data, dict):
+                    # Merge included data with local data (local overrides included?)
+                    # Spec: "Included file data will be included verbatim at the current level of nesting."
+                    # Usually specific (local) overrides generic (included).
+                    # So we take included, update with local.
+                    result = included_data.copy()
+                    result.update(local_data)
+                    return result
+                else:
+                    # If included data is not a dict, we can't merge it into a dict.
+                    # Unless the dict was JUST the include?
+                    if not local_data:
+                        return included_data
+                    # Fallback: return local data (ignore include) or error?
+                    # For now, let's return local_data but maybe warn?
+                    # Or maybe return included_data if local_data is empty?
+                    return local_data
+
+        # Recurse on values if no include or after processing
+        return {k: process_includes(v, base_dir) for k, v in data.items()}
+
+    elif isinstance(data, list):
+        new_list = []
+        for item in data:
+            # Check if item is a dict with ONLY 'include' key
+            if isinstance(item, dict) and "include" in item and len(item) == 1:
+                include_file = item["include"]
+                if isinstance(include_file, str):
+                    filepath = os.path.join(base_dir, include_file)
+                    included_data = load_file_to_dict(filepath, check_extension=False)
+
+                    if isinstance(included_data, list):
+                        new_list.extend(included_data)
+                    else:
+                        new_list.append(included_data)
+                else:
+                    new_list.append(process_includes(item, base_dir))
+            else:
+                new_list.append(process_includes(item, base_dir))
+        return new_list
+
+    else:
+        return data
+
+
+def load_file_to_dict(filename: str, check_extension: bool = True) -> dict:
     # Attempt to strip two levels of file extensions to determine the schema.
     #   Examples: fodo.pals.yaml, fodo.pals.json, ...
     file_noext, extension, file_noext_noext, extension_inner = inspect_file_extensions(
-        filename
+        filename, check_extension=check_extension
     ).values()
 
     # examples: fodo.pals.yaml, fodo.pals.json
@@ -51,6 +134,9 @@ def load_file_to_dict(filename: str) -> dict:
                 f"load_file_to_dict: No support for PALS file {filename} with extension {extension} yet."
             )
 
+    # Process includes
+    pals_data = process_includes(pals_data, base_dir=os.path.dirname(filename))
+
     return pals_data
 
 

src/pals/kinds/Lattice.py

@@ -1,7 +1,8 @@
-from pydantic import model_validator, Field
-from typing import Annotated, List, Literal, Union
+from pydantic import model_validator
+from typing import List, Literal, Union
 
 from .BeamLine import BeamLine
+from .PlaceholderName import PlaceholderName
 from .mixin import BaseElement
 from ..functions import load_file_to_dict, store_dict_to_file
 
@@ -11,7 +12,7 @@ class Lattice(BaseElement):
 
     kind: Literal["Lattice"] = "Lattice"
 
-    branches: List[Annotated[Union[BeamLine], Field(discriminator="kind")]]
+    branches: List[Union[BeamLine, PlaceholderName]]
 
     @model_validator(mode="before")
     @classmethod

tests/test_include.py

@@ -0,0 +1,139 @@
+import pals
+
+
+def test_include(tmp_path):
+    main_file = tmp_path / "main.pals.yaml"
+    root_included_file = tmp_path / "included.pals.yaml"
+    facility_included_file = tmp_path / "facility.pals.yaml"
+    facility_nested_file = tmp_path / "facility_nested.pals.yaml"
+
+    main_content = f"""
+    PALS:
+      include: "{root_included_file.name}"
+      facility:
+        - drift1:
+            kind: Drift
+            length: 0.25
+
+        - include: "{facility_included_file.name}"
+
+        - fodo_cell:
+            kind: BeamLine
+            line:
+              - drift1
+              - quad1
+              - drift2
+              - quad2
+              - drift1
+    
+        - fodo_lattice:
+            kind: Lattice
+            branches:
+              - fodo_cell
+    
+        - use: fodo_lattice
+    """
+
+    root_included_content = """
+    author: "Some One <name@email.com>"
+    version: 1.0
+    """
+
+    facility_included_content = f"""
+    - quad1:
+        kind: Quadrupole
+        MagneticMultipoleP:
+          Bn1: 1.0
+        length: 1.0
+
+    - drift2:
+        kind: Drift
+        length: 0.5
+    
+    - include: "{facility_nested_file.name}"
+    """
+
+    facility_nested_content = """
+        - quad2:
+            kind: Quadrupole
+            MagneticMultipoleP:
+              Bn1: -1.0
+            length: 1.0
+        """
+
+    main_file.write_text(main_content)
+    root_included_file.write_text(root_included_content)
+    facility_included_file.write_text(facility_included_content)
+    facility_nested_file.write_text(facility_nested_content)
+
+    data = pals.Lattice.from_file(main_file)
+
+    assert data["PALS"]["version"] == 1.0
+    assert data["PALS"]["other"] == "value"
+    assert data["PALS"]["author"] == "Some One <name@email.com>"
+    assert "include" not in data["PALS"]
+
+
+def test_nested_include(tmp_path):
+    root_file = tmp_path / "root.pals.yaml"
+    middle_file = tmp_path / "middle.pals.yaml"
+    leaf_file = tmp_path / "leaf.pals.yaml"
+
+    root_content = f"""
+    root:
+      include: "{middle_file.name}"
+    """
+
+    middle_content = f"""
+    middle: val
+    include: "{leaf_file.name}"
+    """
+
+    leaf_content = """
+    leaf: val
+    """
+
+    root_file.write_text(root_content)
+    middle_file.write_text(middle_content)
+    leaf_file.write_text(leaf_content)
+
+    data = pals.functions.load_file_to_dict(str(root_file))
+
+    assert data["root"]["middle"] == "val"
+    assert data["root"]["leaf"] == "val"
+    assert "include" not in data["root"]
+
+
+def test_include_list_into_dict_conversion(tmp_path):
+    # This tests the spec example where a list of properties is included into a dict (element)
+    main_file = tmp_path / "element.pals.yaml"
+    params_file = tmp_path / "params.pals.yaml"
+
+    main_content = f"""
+    element:
+      kind: Quadrupole
+      include: "{params_file.name}"
+    """
+
+    # params file content is a list of single-key dicts
+    params_content = """
+    - MagneticMultipoleP:
+      - Kn3L: 0.3
+    - ApertureP:
+        x_limits: [-0.1, 0.1]
+    """
+
+    main_file.write_text(main_content)
+    params_file.write_text(params_content)
+
+    data = pals.functions.load_file_to_dict(str(main_file))
+
+    elem = data["element"]
+    assert elem["kind"] == "Quadrupole"
+
+    # Check if keys are correctly merged from the list
+    assert "MagneticMultipoleP" in elem
+    assert elem["MagneticMultipoleP"] == [{"Kn3L": 0.3}]
+
+    assert "ApertureP" in elem
+    assert elem["ApertureP"] == {"x_limits": [-0.1, 0.1]}