typehint for 3.8

lint lint lint lint lint lint lint lint
materialsproject · Jan 31, 2023 · 2fa05bb · 2fa05bb
1 parent c84e913
commit 2fa05bb
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 10 deletions.
diff --git a/src/atomate2/common/analysis/defects/jobs.py b/src/atomate2/common/analysis/defects/jobs.py
@@ -26,8 +26,14 @@
 
 logger = logging.getLogger(__name__)
 
-
-logger = logging.getLogger(__name__)
+__all__ = [
+    "get_charged_structures",
+    "spawn_energy_curve_calcs",
+    "get_ccd_documents",
+    "get_supercell_from_prv_calc",
+    "bulk_supercell_calculation",
+    "spawn_defect_calcs",
+]
 
 
 class CCDInput(BaseModel):

diff --git a/src/atomate2/common/schemas/defects.py b/src/atomate2/common/schemas/defects.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Callable, List, Optional, Tuple, Type
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 from pydantic import BaseModel, Field
@@ -12,8 +12,6 @@
 
 logger = logging.getLogger(__name__)
 
-__all__ = ["CCDDocument", "FormationEnergyDiagramDocument"]
-
 
 class FormationEnergyDiagramDocument(BaseModel):
     """A document for storing a formation energy diagram.
@@ -55,11 +53,11 @@ class FormationEnergyDiagramDocument(BaseModel):
         None, description="The directory name of the pristine supercell calculation."
     )
 
-    defect_sc_dirs: dict[int, str] = Field(
+    defect_sc_dirs: Dict[int, str] = Field(
         None, description="The directory names of the charged defect calculations."
     )
 
-    dielectric: float | list[list[float]] = Field(
+    dielectric: Union[float, List[List[float]]] = Field(
         None,
         description="The dielectric constant or tensor, can be used to compute corrections.",
     )
@@ -88,7 +86,7 @@ def from_FormationEnergyDiagram(
 
     def as_FormationEnergyDiagram(
         self, pd_entries: Optional[List[ComputedEntry]] = None
-    ) -> "FormationEnergyDiagram":
+    ) -> FormationEnergyDiagram:
         """Create a `FormationEnergyDiagram` object from the document.
 
         Since the `pd_entries` field is optional, this method allows the user
@@ -235,7 +233,7 @@ def get_ent(struct, energy, dir_name, uuid):
 
     @classmethod
     def from_entries(
-        cls: Type["CCDDocument"],
+        cls,
         entries1: List[ComputedStructureEntry],
         entries2: List[ComputedStructureEntry],
         relaxed_uuid1: Optional[str] = None,
@@ -257,7 +255,7 @@ def from_entries(
 
         """
 
-        def find_entry(entries, uuid) -> Tuple[int, ComputedStructureEntry]:
+        def find_entry(entries, uuid) -> tuple[int, ComputedStructureEntry]:
             """Find the entry with the given UUID."""
             for itr, entry in enumerate(entries):
                 if entry.data["uuid"] == uuid:

diff --git a/src/atomate2/vasp/jobs/defect.py b/src/atomate2/vasp/jobs/defect.py
@@ -8,6 +8,14 @@
 from pymatgen.io.vasp import Incar
 from pymatgen.io.vasp.outputs import WSWQ
 
+from atomate2.common.analysis.defects.jobs import (
+    bulk_supercell_calculation,
+    get_ccd_documents,
+    get_charged_structures,
+    get_supercell_from_prv_calc,
+    spawn_defect_calcs,
+    spawn_energy_curve_calcs,
+)
 from atomate2.common.files import copy_files, gunzip_files, gzip_files, rename_files
 from atomate2.utils.file_client import FileClient
 from atomate2.utils.path import strip_hostname
@@ -17,6 +25,17 @@
 
 logger = logging.getLogger(__name__)
 
+__all__ = [
+    "bulk_supercell_calculation",
+    "calculate_finite_diff",
+    "get_ccd_documents",
+    "get_charged_structures",
+    "get_supercell_from_prv_calc",
+    "spawn_defect_calcs",
+    "spawn_energy_curve_calcs",
+]
+# sort the list above
+
 
 @job(data=WSWQ, output_schema=FiniteDifferenceDocument)
 def calculate_finite_diff(