added first iteration of ess.TraceWin.project

ess/TraceWin.py

@@ -1492,3 +1492,131 @@ class field_map:
         data = self.map.reshape(totmapshape)
         for j in data:
             fout.write("{}\n".format(j))
+
+
+class project:
+    def __init__(self, project_fname=None, settings_fname=None):
+        """
+        Read and modify project files
+
+        Example::
+            p = project('SPK.ini')
+            for diff in p.compare_to('MEBT.ini'):
+                print(diff)
+            p.set('main:beam1_energy', 89e6)
+            p.save()
+        """
+        import json
+        import pkg_resources
+
+        if settings_fname is None:
+            self._refdict = json.loads(pkg_resources.resource_string(__name__, "data/tw_project_file_reverse_engineered.json"))
+        else:
+            self._refdict = json.loads(open(self._settings_fname, "r").read())
+
+        self._settings_fname = settings_fname
+        self._project_fname = project_fname
+
+        self._dict = {}
+        if self._project_fname is not None:
+            self._read_settings()
+
+    def _read_settings(self):
+        import struct
+        import textwrap
+
+        with open(self._project_fname, "rb") as f:
+            hexlist = textwrap.wrap(f.read().hex(), 2)
+        for key in self._refdict:
+            o = self._refdict[key]
+            current = "".join(hexlist[o[0] : o[0] + o[2]])
+            if o[1] == "bool":
+                if current == o[3]:
+                    self._dict[key] = False
+                elif current == o[4]:
+                    self._dict[key] = True
+                else:
+                    raise ValueError(f"Wrong setting found for {key}, found {current}, expected one of {o[3:]}")
+            elif o[1] in ["d", "f", "i"]:
+                self._dict[key] = struct.unpack(o[1], b"".fromhex(current))[0]
+
+    def print_settings(self, settings=None):
+        """
+        Print the settings given, or all by default
+
+        :param settings: List of the settings to print
+        """
+        if settings is None:
+            settings = self.keys()
+
+        for setting in settings:
+            print(setting, self._dict[setting])
+
+    def keys(self):
+        return self._dict.keys()
+
+    def get(self, parameter):
+        """
+        Get the setting of the parameter
+        """
+        return self._dict[parameter]
+
+    def set(self, parameter, value):
+        """
+        Set the new value for parameter
+        """
+        current = self.get(parameter)
+        if isinstance(current, bool):
+            if not isinstance(value, bool):
+                raise ValueError(f"{parameter} should be True or False")
+        elif isinstance(current, (float, int)):
+            if not isinstance(value, (float, int)):
+                raise ValueError(f"{parameter} should be a number")
+        self._dict[parameter] = value
+
+    def save(self, fname=None):
+        """
+        Save the project file
+
+        If fname not given, overwrite original file
+        """
+        import struct
+        from textwrap import wrap
+
+        if fname is None:
+            fname = self._project_fname
+
+        with open(self._project_fname, "rb") as f:
+            hexlist = wrap(f.read().hex(), 2)
+
+        for key in self._dict:
+            o = self._refdict[key]
+            v = self._dict[key]
+            if isinstance(v, bool):
+                if v:
+                    v = [o[-1]]
+                else:
+                    v = [o[-2]]
+            else:
+                v = wrap(struct.pack(o[1], v).hex(), 2)
+            for i in range(len(v)):
+                hexlist[o[0] + i] = v[i]
+
+        open(fname, "wb").write(bytes.fromhex("".join(hexlist)))
+
+    def compare_to(self, other):
+        """
+        Compare the settings of this file to a
+        different project file
+
+        :param other: project object, or file path to other project file
+        """
+        diffs = []
+
+        if isinstance(other, str):
+            other = project(other)
+
+        for key in self.keys():
+            if self.get(key) != other.get(key):
+                diffs.append([key, self.get(key), other.get(key)])
+        return diffs

ess/data/tw_project_file_reverse_engineered.json

+{
+  "partran": [
+    109,
+    "bool",
+    1,
+    "00",
+    "01"
+  ],
+  "toutatis": [
+    135,
+    "bool",
+    1,
+    "00",
+    "01"
+  ],
+  "match:input_matched_beam": [
+    128,
+    "bool",
+    1,
+    "00",
+    "01"
+  ],
+  "match:family_twiss": [
+    113,
+    "bool",
+    1,
+    "00",
+    "01"
+  ],
+  "match:diag": [
+    118,
+    "bool",
+    1,
+    "00",
+    "01"
+  ],
+  "match:limit_criterion": [
+    12260,
+    "d",
+    8
+  ],
+  "main:beam1_npart": [
+    11984,
+    "i",
+    4
+  ],
+  "main:beam2_npart": [
+    11988,
+    "i",
+    4
+  ],
+  "main:beam1_energy": [
+    12100,
+    "d",
+    8
+  ],
+  "main:beam1_freq": [
+    12068,
+    "d",
+    8
+  ],
+  "main:beam1_curr": [
+    12084,
+    "d",
+    8
+  ],
+  "main:beam1_dcycle": [
+    12648,
+    "d",
+    8
+  ],
+  "main:beam1_emit_xxp": [
+    12116,
+    "d",
+    8
+  ],
+  "main:beam1_emit_yyp": [
+    12132,
+    "d",
+    8
+  ],
+  "main:beam1_rms_energy_spread": [
+    22752,
+    "d",
+    8
+  ],
+  "main:beam1_twiss_alphaz": [
+    12572,
+    "d",
+    8
+  ],
+  "main:beam1_twiss_betaz": [
+    12580,
+    "d",
+    8
+  ]
+}

setup.py

@@ -26,5 +26,6 @@ setup(
         "Programming Language :: Python :: 3.8",
     ],
     packages=["ess"],
+    package_data={"ess": ["data/*.json"]},
     scripts=["scripts/twcli", "scripts/tracewin_errorstudy", "scripts/ibsimu2tw"],
 )