Fix bug in regex for perturbed atoms

aiidateam · Dec 5, 2023 · 7aa81fd · 7aa81fd
1 parent 4e28e25
commit 7aa81fd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/aiida_quantumespresso_hp/utils/general.py b/src/aiida_quantumespresso_hp/utils/general.py
@@ -28,7 +28,7 @@ def is_perturb_only_atom(parameters: dict) -> int | None:
  match = None # making sure that if the dictionary is empty we don't raise an `UnboundLocalError`
 
  for key in parameters.keys():
- match = re.search(r'perturb_only_atom.*([0-9]).*', key)
+ match = re.search(r'perturb_only_atom.*?(\d+).*', key)
  if match:
  if not parameters[key]: # also the key must be `True`
  match = None # making sure to have `None`

diff --git a/tests/utils/test_general.py b/tests/utils/test_general.py
@@ -26,5 +26,8 @@ def test_is_perturb_only_atom():
  parameters = {'perturb_only_atom(1)': True}
  assert is_perturb_only_atom(parameters) == 1
 
+ parameters = {'perturb_only_atom(20)': True}
+ assert is_perturb_only_atom(parameters) == 20
+
  parameters = {'perturb_only_atom(1)': False}
  assert is_perturb_only_atom(parameters) is None