added merge sets

Author: mcocdawc

src/quemb/molbe/chemfrag.py

@@ -4,6 +4,8 @@
 from chemcoord import Cartesian
 from pyscf.gto import Mole
 
+from quemb.shared.typing import T
+
 AtomIdx = NewType("AtomIdx", int)
 
 AOIdx = NewType("AOIdx", int)
@@ -38,6 +40,10 @@
 ContainedCenterIdx = NewType("ContainedCenterIdx", dict[CenterIdx, set[CenterIdx]])
 
 
+def merge_sets(*sets: set[T]) -> set[T]:
+    return set().union(*sets)
+
+
 @define
 class FragmentedMolecule:
     atom_per_frag: Final[AtomPerFrag]
@@ -155,7 +161,7 @@ def AO_indices_of_fragment(fragment: set[AtomIdx]) -> dict[AtomIdx, set[AOIdx]]:
 def get_fsites(mol: Mole, fragments: AtomPerFrag) -> AOPerFrag:
     return AOPerFrag(
         {
-            i_center: set().union(*i_fragment.values())
+            i_center: merge_sets(*i_fragment.values())
             for i_center, i_fragment in get_fs(mol, fragments).items()
         }
     )

src/quemb/shared/typing.py

@@ -48,3 +48,7 @@
 PathLike: TypeAlias = str | os.PathLike
 #: Type annotation for dictionaries holding keyword arguments.
 KwargDict: TypeAlias = Dict[str, Any]
+
+
+#: A generic type variable, without any constraints.
+T = TypeVar("T")