Add support for complex entries

dftbplus · Mar 13, 2023 · 0d97116 · 0d97116
1 parent 2176653
commit 0d97116
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 6 deletions.
diff --git a/src/hsd/dict.py b/src/hsd/dict.py
@@ -13,24 +13,49 @@
     QUOTING_CHARS, SPECIAL_CHARS
 from hsd.eventhandler import HsdEventHandler, HsdEventPrinter
 
-_ItemType = Union[float, int, bool, str]
+_ItemType = Union[float, complex, int, bool, str]
 
 _DataType = Union[_ItemType, List[_ItemType]]
 
 # Pattern to transform HSD string values into actual Python data types
 _TOKEN_PATTERN = re.compile(r"""
 # Integer
-(?:\s*(?:^|(?<=\s))(?P<int>[+-]?[0-9]+)(?:\s*$|\s+))
+(?:
+    \s*(?:^|(?<=\s))
+    (?P<int>[+-]?[0-9]+)
+    (?:\s*$|\s+)
+)
 |
 # Floating point
-(?:\s*(?:^|(?<=\s))
-(?P<float>[-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)(?:$|(?=\s+)))
+(?:
+    \s*(?:^|(?<=\s))
+    (?P<float>[-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)
+    (?:$|(?=\s+))
+)
+|
+# Complex value
+(?:
+    \s*(?:^|(?<=\s))
+    (?P<complex>
+        [-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?i
+        |
+        [-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?[-+][0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?i
+    )
+    (?:$|(?=\s+))
+)
 |
 # Logical (Yes/No)
-(?:\s*(?:^|(?<=\s))(?P<logical>[Yy][Ee][Ss]|[Nn][Oo])(?:$|(?=\s+)))
+(?:
+    \s*(?:^|(?<=\s))
+    (?P<logical>[Yy][Ee][Ss]|[Nn][Oo])
+    (?:$|(?=\s+))
+)
 |
 # Quoted string
-(?:\s*(?:(?P<qstr>(?P<quote>['"]).*?(?P=quote))
+(?:
+    \s*
+    (?:(?P<qstr>(?P<quote>['"]).*?(?P=quote)
+)
 |
 # Unquoted string
 (?P<str>.+?))(?:$|\s+))
@@ -148,6 +173,9 @@ def _text_to_data(self, txt: str) -> _DataType:
                     linedata.append(int(match.group("int")))
                 elif match.group("float"):
                     linedata.append(float(match.group("float")))
+                elif match.group("complex"):
+                    pycmplx = match.group("complex").replace("i", "j", 1)
+                    linedata.append(complex(pycmplx))
                 elif match.group("logical"):
                     lowlog = match.group("logical").lower()
                     linedata.append(lowlog == "yes")
@@ -246,6 +274,8 @@ def _item_to_hsd(item):
         return "Yes" if item else "No"
     if isinstance(item, (int, float)):
         return str(item)
+    if isinstance(item, complex):
+        return f"{item.real}{item.imag:+}i"
     if isinstance(item, str):
         return _str_to_hsd(item)
     msg = "Data type {} can not be converted to HSD string"\

diff --git a/test/test_dict.py b/test/test_dict.py
@@ -78,6 +78,15 @@
             }
         )
     ),
+    (
+        "Complex entries", (
+            "ComplexNums = 4.0+9.0i 2.1e-12-4.5e-12i 0.32+0.45i -0.01-1.0i\n",
+            {
+                "ComplexNums.hsdattrib": {_HSD_EQUAL: True, _HSD_LINE: 0},
+                "ComplexNums": [4.0+9.0j, 2.1e-12-4.5e-12j, 3.2e-1+4.5e-1j, -0.01-1.0j]
+            }
+        )
+    ),
     (
         "Duplicate node", (
             "a {\n  b = 1\n}\na {\n  b = 2\n}\n",