From 0e1ca7a46532893d0caf39f4c808acbf0687edf5 Mon Sep 17 00:00:00 2001
From: Ivan Usov <ivan.usov@psi.ch>
Date: Thu, 23 Jul 2020 20:43:37 +0200
Subject: [PATCH] Handle missing tags in crystal element

---
 pyzebra/anatric.py | 48 +++++++++++++++++++++++++---------------------
 1 file changed, 26 insertions(+), 22 deletions(-)

diff --git a/pyzebra/anatric.py b/pyzebra/anatric.py
index 5ed414c..5adb9ad 100644
--- a/pyzebra/anatric.py
+++ b/pyzebra/anatric.py
@@ -42,6 +42,22 @@ class AnatricConfig:
     def save_as(self, filename):
         self._tree.write(filename)
 
+    def _get_attr(self, name, tag, attr):
+        elem = self._tree.find(name).find(tag)
+        if elem is None:
+            return None
+        return elem.attrib[attr]
+
+    def _set_attr(self, name, tag, attr, value):
+        tree_elem = self._tree.find(name)
+        elem = tree_elem.find(tag)
+        if elem is None:
+            new_elem = ET.Element(tag, attrib={attr: value})
+            new_elem.tail = "\n"
+            tree_elem.append(new_elem)
+        else:
+            elem.attrib[attr] = value
+
     @property
     def logfile(self):
         return self._tree.find("logfile").attrib["file"]
@@ -113,55 +129,43 @@ class AnatricConfig:
 
     @property
     def crystal_sample(self):
-        return self._tree.find("crystal").find("Sample").attrib["name"]
+        return self._get_attr("crystal", "Sample", "name")
 
     @crystal_sample.setter
     def crystal_sample(self, value):
-        self._tree.find("crystal").find("Sample").attrib["name"] = value
+        self._set_attr("crystal", "Sample", "name", value)
 
     @property
     def crystal_lambda(self):
-        elem = self._tree.find("crystal").find("lambda")
-        if elem is not None:
-            return elem.attrib["value"]
-        return None
+        return self._get_attr("crystal", "lambda", "value")
 
     @crystal_lambda.setter
     def crystal_lambda(self, value):
-        self._tree.find("crystal").find("lambda").attrib["value"] = value
+        self._set_attr("crystal", "lambda", "value", value)
 
     @property
     def crystal_zeroOM(self):
-        elem = self._tree.find("crystal").find("zeroOM")
-        if elem is not None:
-            return elem.attrib["value"]
-        return None
+        return self._get_attr("crystal", "zeroOM", "value")
 
     @crystal_zeroOM.setter
     def crystal_zeroOM(self, value):
-        self._tree.find("crystal").find("zeroOM").attrib["value"] = value
+        self._set_attr("crystal", "zeroOM", "value", value)
 
     @property
     def crystal_zeroSTT(self):
-        elem = self._tree.find("crystal").find("zeroSTT")
-        if elem is not None:
-            return elem.attrib["value"]
-        return None
+        return self._get_attr("crystal", "zeroSTT", "value")
 
     @crystal_zeroSTT.setter
     def crystal_zeroSTT(self, value):
-        self._tree.find("crystal").find("zeroSTT").attrib["value"] = value
+        self._set_attr("crystal", "zeroSTT", "value", value)
 
     @property
     def crystal_zeroCHI(self):
-        elem = self._tree.find("crystal").find("zeroCHI")
-        if elem is not None:
-            return elem.attrib["value"]
-        return None
+        return self._get_attr("crystal", "zeroCHI", "value")
 
     @crystal_zeroCHI.setter
     def crystal_zeroCHI(self, value):
-        self._tree.find("crystal").find("zeroCHI").attrib["value"] = value
+        self._set_attr("crystal", "zeroCHI", "value", value)
 
     @property
     def crystal_UB(self):