From d532d7f63ad3932e52ada7725a2a1c4aab4f9d09 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 14 Nov 2024 16:29:16 -0800 Subject: [PATCH] update type hinting / docstr --- src/atomate2/common/jobs/electrode.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/atomate2/common/jobs/electrode.py b/src/atomate2/common/jobs/electrode.py index 924b59e85..0e19de873 100644 --- a/src/atomate2/common/jobs/electrode.py +++ b/src/atomate2/common/jobs/electrode.py @@ -21,6 +21,7 @@ from pymatgen.analysis.structure_matcher import StructureMatcher from pymatgen.core import Structure from pymatgen.entries.computed_entries import ComputedEntry + from pymatgen.io.common import VolumetricData logger = logging.getLogger(__name__) @@ -203,7 +204,7 @@ def get_insertion_electrode_doc( @job def get_inserted_structures( prev_dir: Path | str, - get_charge_density: Callable, + get_charge_density: Callable[[str | Path], VolumetricData], inserted_species: ElementLike, insertions_per_step: int = 4, charge_insertion_generator: ChargeInterstitialGenerator | None = None, @@ -213,7 +214,7 @@ def get_inserted_structures( Parameters ---------- prev_dir: The previous directory where the static calculation was performed. - get_charge_density: A function to get the charge density from a task document. + get_charge_density: A function to get the charge density from a run directory. inserted_species: The species to insert. insertions_per_step: The maximum number of ion insertion sites to attempt. charge_insertion_generator: The charge insertion generator to use,