pals-python/tests/test_include.py at ffd096b53be6589a76666bf7e897931145c5e175 · pals-project/pals-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
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]}