Refactored a utility function

Author: RaphaelRobidas

ccinput/packages/xtb.py

@@ -1,7 +1,7 @@
 import os
 
 from ccinput.constants import CalcType, ATOMIC_NUMBER, LOWERCASE_ATOMIC_SYMBOLS
-from ccinput.utilities import get_solvent, get_method
+from ccinput.utilities import get_solvent, get_method, compress_indices
 from ccinput.exceptions import InvalidParameter, ImpossibleCalculation
 
 
@@ -128,33 +128,6 @@ def handle_constraints_scan(self):
                 if cmd.scan:
                     self.input_file += f"{counter+1}: {cmd.start_d:.2f}, {cmd.end_d:.2f}, {cmd.num_steps}\n"
 
-    def compress_indices(self, arr):
-        comp = []
-
-        def add_to_str(curr):
-            if len(curr) == 0:
-                return ""
-            elif len(curr) == 1:
-                return f"{curr[0]}"
-            else:
-                return f"{curr[0]}-{curr[-1]}"
-
-        _arr = sorted(set(arr))
-        curr_atoms = []
-
-        for a in _arr:
-            if len(curr_atoms) == 0:
-                curr_atoms.append(a)
-            else:
-                if a == curr_atoms[-1] + 1:
-                    curr_atoms.append(a)
-                else:
-                    comp.append(add_to_str(curr_atoms))
-                    curr_atoms = [a]
-
-        comp.append(add_to_str(curr_atoms))
-        return ",".join(comp)
-
     def handle_constraints_crest(self):
         if len(self.calc.constraints) == 0:
             raise InvalidParameter("No constraint in constrained optimisation mode")
@@ -171,15 +144,15 @@ def handle_constraints_crest(self):
             self.input_file += cmd.to_xtb()
             constr_atoms += cmd.ids
 
-        self.input_file += f"atoms: {self.compress_indices(constr_atoms)}\n"
+        self.input_file += f"atoms: {compress_indices(constr_atoms)}\n"
 
         mtd_atoms = list(range(1, num_atoms))
         for a in constr_atoms:
             if int(a) in mtd_atoms:
                 mtd_atoms.remove(int(a))
 
         self.input_file += "$metadyn\n"
-        self.input_file += f"atoms: {self.compress_indices(mtd_atoms)}\n"
+        self.input_file += f"atoms: {compress_indices(mtd_atoms)}\n"
 
     def handle_specifications(self):
         accuracy = -1

ccinput/utilities.py

@@ -2,8 +2,25 @@
 import string
 import numpy as np
 
-from ccinput.constants import *
-from ccinput.exceptions import *
+from ccinput.constants import (
+    THEORY_LEVELS,
+    SOFTWARE_METHODS,
+    SOFTWARE_BASIS_SETS,
+    SOFTWARE_SOLVENTS,
+    SYN_TYPES,
+    SYN_METHODS,
+    SYN_SOFTWARE,
+    SYN_BASIS_SETS,
+    SYN_SOLVENTS,
+    ATOMIC_NUMBER,
+    ATOMIC_SYMBOL,
+    LOWERCASE_ATOMIC_SYMBOLS,
+    FUNCTIONALS_WITH_DISPERSION_PARAMETERS,
+    BASIS_SET_EXCHANGE_KEY,
+    EXCHANGE_FUNCTIONALS,
+    CORRELATION_FUNCTIONALS,
+)
+from ccinput.exceptions import InvalidParameter, InvalidXYZ
 
 MEMORY_FACTORS = {
     "m": 1,
@@ -520,3 +537,31 @@ def parse_specifications(specs, add_option_fn, condense=True):
             add_option_fn(key, option.replace("&", " "))
         else:
             add_option_fn(spec, "")
+
+
+def compress_indices(arr):
+    comp = []
+
+    def add_to_str(curr):
+        if len(curr) == 0:
+            return ""
+        elif len(curr) == 1:
+            return f"{curr[0]}"
+        else:
+            return f"{curr[0]}-{curr[-1]}"
+
+    _arr = sorted(set(arr))
+    curr_atoms = []
+
+    for a in _arr:
+        if len(curr_atoms) == 0:
+            curr_atoms.append(a)
+        else:
+            if a == curr_atoms[-1] + 1:
+                curr_atoms.append(a)
+            else:
+                comp.append(add_to_str(curr_atoms))
+                curr_atoms = [a]
+
+    comp.append(add_to_str(curr_atoms))
+    return ",".join(comp)