Improve type hints in io.ase (#4556)

* type hinting in io/ase

* type hinting in io/ase

Author: tomdemeyere

src/pymatgen/io/ase.py

@@ -8,7 +8,7 @@
 import warnings
 from copy import deepcopy
 from importlib.metadata import PackageNotFoundError
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Literal, TypeVar, overload
 
 import numpy as np
 from monty.json import MontyDecoder, MSONable, jsanitize
@@ -48,8 +48,8 @@ def __init__(self, *args, **kwargs):
 __email__ = "[email protected]"
 __date__ = "Mar 8, 2012"
 
-StructT = TypeVar("StructT", bound=IStructure | IMolecule | Structure | Molecule)
-MolT = TypeVar("MolT", bound=IMolecule)
+IMoleculeT = TypeVar("IMoleculeT", bound=IMolecule)
+StructOrMolT = TypeVar("StructOrMolT", bound=Structure | Molecule)
 
 
 class MSONAtoms(Atoms, MSONable):
@@ -86,6 +86,30 @@ def from_dict(cls, dct: dict[str, Any]) -> Self:
 class AseAtomsAdaptor:
     """Adaptor serves as a bridge between ASE Atoms and pymatgen objects."""
 
+    @overload
+    @staticmethod
+    def get_atoms(
+        structure: SiteCollection,
+        msonable: Literal[True] = ...,
+        **kwargs: Any,
+    ) -> MSONAtoms: ...
+
+    @overload
+    @staticmethod
+    def get_atoms(
+        structure: SiteCollection,
+        msonable: Literal[False],
+        **kwargs: Any,
+    ) -> Atoms: ...
+
+    @overload
+    @staticmethod
+    def get_atoms(
+        structure: SiteCollection,
+        msonable: bool = True,
+        **kwargs: Any,
+    ) -> MSONAtoms | Atoms: ...
+
     @staticmethod
     def get_atoms(
         structure: SiteCollection,
@@ -239,9 +263,9 @@ def get_atoms(
     @staticmethod
     def get_structure(
         atoms: Atoms,
-        cls=Structure,
+        cls: type[StructOrMolT] = Structure,
         **cls_kwargs,
-    ) -> Structure | Molecule:
+    ) -> StructOrMolT:
         """Get pymatgen structure from ASE Atoms.
 
         Args:
@@ -392,7 +416,7 @@ def get_structure(
         return structure
 
     @staticmethod
-    def get_molecule(atoms: Atoms, cls: type[MolT] = Molecule, **cls_kwargs) -> Molecule | IMolecule:  # type:ignore[assignment]
+    def get_molecule(atoms: Atoms, cls: type[IMoleculeT] = Molecule, **cls_kwargs) -> IMoleculeT:
         """Get pymatgen molecule from ASE Atoms.
 
         Args:
@@ -401,7 +425,7 @@ def get_molecule(atoms: Atoms, cls: type[MolT] = Molecule, **cls_kwargs) -> Mole
             **cls_kwargs: Any additional kwargs to pass to the cls constructor
 
         Returns:
-            (I)Molecule: Equivalent pymatgen (I)Molecule
+            MolT: Equivalent pymatgen (I)Molecule
         """
         molecule = AseAtomsAdaptor.get_structure(atoms, cls=cls, **cls_kwargs)