Skip to content
Draft
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 90 additions & 4 deletions src/pals/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand All @@ -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
Expand All @@ -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


Expand Down
7 changes: 4 additions & 3 deletions src/pals/kinds/Lattice.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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
Expand Down
139 changes: 139 additions & 0 deletions tests/test_include.py
Original file line number Diff line number Diff line change
@@ -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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What was the intention here? Why does this go through Lattice.from_file as opposed to going through functions.load_file_to_dict as below (which would work here too - tested locally)? Could you comment on whether you wanted to expand the Lattice.from_file function but haven't gotten to implement the feature in the PR yet? I think I'm missing some context, inteded use case, and current PR progress, to start working on completing this PR.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ax3l

Do you have any info to share about the questions in my previous comment?

@ax3l ax3l Jun 22, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I thought we discussed this in person 2 weeks ago :)

from_file is the primary user-facing API users shall use, so this exercises that API.
Using both here and below just spreads the API surface coverage with minimal duplication.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, thanks. Let me know if the changes I pushed in 1b437b4 accomplish what you had in mind.


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]}
Loading