Fix types

edan-bainglass · Aug 12, 2024 · 3a1fea0 · 3a1fea0
1 parent a41c5fb
commit 3a1fea0
Show file tree

Hide file tree

Showing 8 changed files with 45 additions and 67 deletions.
diff --git a/src/aiida_quantum_transport/calculations/base.py b/src/aiida_quantum_transport/calculations/base.py
@@ -1,20 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from aiida.engine import CalcJob
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class BaseCalculation(CalcJob):
     """docstring"""
 
     _default_parser_name = ""
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)

diff --git a/src/aiida_quantum_transport/calculations/current.py b/src/aiida_quantum_transport/calculations/current.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class CurrentCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.current"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -114,12 +110,14 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             transmission_folder_path,
         ]
 
-        hybridization_data: orm.RemoteData = (
-            self.inputs.hybridization.remote_results_folder
-        )
-        transmission_data: orm.RemoteData = (
-            self.inputs.transmission.remote_results_folder
-        )
+        hybridization_data = self.inputs.hybridization.remote_results_folder
+        transmission_data = self.inputs.transmission.remote_results_folder
+
+        for rd in (hybridization_data, transmission_data):
+            if not isinstance(rd, orm.RemoteData):
+                raise ValueError(f"Expected `RemoteData` instance; got `{type(rd)}`")
+            if rd.computer is None:
+                raise ValueError("Missing `Computer` node")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/dft.py b/src/aiida_quantum_transport/calculations/dft.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class DFTCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.dft"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)

diff --git a/src/aiida_quantum_transport/calculations/dmft.py b/src/aiida_quantum_transport/calculations/dmft.py
@@ -2,7 +2,6 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import numpy as np
 from aiida import orm
@@ -11,17 +10,14 @@
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class DMFTCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.dmft"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -158,12 +154,8 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         matsubara_hybridization_filepath = (
             precomputed_input_dir / "matsubara_hybridization.bin"
         ).as_posix()
-        hamiltonian_filepath = (
-            precomputed_input_dir / "hamiltonian.npy"
-        ).as_posix()
-        occupancies_filepath = (
-            precomputed_input_dir / "occupancies.npy"
-        ).as_posix()
+        hamiltonian_filepath = (precomputed_input_dir / "hamiltonian.npy").as_posix()
+        occupancies_filepath = (precomputed_input_dir / "occupancies.npy").as_posix()
         scattering_region_filename = "scatt.npy"
         scattering_region_filepath = (
             precomputed_input_dir / scattering_region_filename
@@ -200,9 +192,14 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             occupancies_filepath,
         ]
 
-        hybridization_data: orm.RemoteData = (
-            self.inputs.hybridization.remote_results_folder
-        )
+        hybridization_data = self.inputs.hybridization.remote_results_folder
+
+        if not isinstance(hybridization_data, orm.RemoteData):
+            raise ValueError(
+                f"Expected `RemoteData` instance; got `{type(hybridization_data)}`"
+            )
+        if hybridization_data.computer is None:
+            raise ValueError("Missing `Computer` node")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/greens.py b/src/aiida_quantum_transport/calculations/greens.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class GreensFunctionParametersCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.greens"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -132,12 +128,8 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
 
         precomputed_input_dir = input_dir / "precomputed"
         (temp_dir / precomputed_input_dir).mkdir()
-        leads_hamiltonian_filepath = (
-            precomputed_input_dir / "hs_leads.npy"
-        ).as_posix()
-        los_hamiltonian_filepath = (
-            precomputed_input_dir / "hs_los.npy"
-        ).as_posix()
+        leads_hamiltonian_filepath = (precomputed_input_dir / "hs_leads.npy").as_posix()
+        los_hamiltonian_filepath = (precomputed_input_dir / "hs_los.npy").as_posix()
 
         codeinfo = CodeInfo()
         codeinfo.code_uuid = self.inputs.code.uuid
@@ -156,8 +148,14 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             los_hamiltonian_filepath,
         ]
 
-        leads_data: orm.RemoteData = self.inputs.leads.remote_results_folder
-        los_data: orm.RemoteData = self.inputs.los.remote_results_folder
+        leads_data = self.inputs.leads.remote_results_folder
+        los_data = self.inputs.los.remote_results_folder
+
+        for rd in (leads_data, los_data):
+            if not isinstance(rd, orm.RemoteData):
+                raise ValueError(f"Expected `RemoteData` instance; got `{type(rd)}`")
+            if rd.computer is None:
+                raise ValueError("Missing `Computer` node")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/hybridize.py b/src/aiida_quantum_transport/calculations/hybridize.py
@@ -2,25 +2,21 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 from aiida import orm
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class HybridizationCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.hybridize"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)
@@ -152,9 +148,7 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
 
         precomputed_input_dir = input_dir / "precomputed"
         (temp_dir / precomputed_input_dir).mkdir()
-        los_indices_filepath = (
-            precomputed_input_dir / "los_indices.npy"
-        ).as_posix()
+        los_indices_filepath = (precomputed_input_dir / "los_indices.npy").as_posix()
         hamiltonian_ii_filepath = (
             precomputed_input_dir / "hamiltonian_ii.pkl"
         ).as_posix()
@@ -180,10 +174,14 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
             self_energies_filepath,
         ]
 
-        los_data: orm.RemoteData = self.inputs.los.remote_results_folder
-        greens_function_data: orm.RemoteData = (
-            self.inputs.greens_function.remote_results_folder
-        )
+        los_data = self.inputs.los.remote_results_folder
+        greens_function_data = self.inputs.greens_function.remote_results_folder
+
+        for rd in (los_data, greens_function_data):
+            if not isinstance(rd, orm.RemoteData):
+                raise ValueError(f"Expected `RemoteData` instance; got `{type(rd)}`")
+            if rd.computer is None:
+                raise ValueError("Missing `Computer` node")
 
         calcinfo = CalcInfo()
         calcinfo.codes_info = [codeinfo]

diff --git a/src/aiida_quantum_transport/calculations/localize.py b/src/aiida_quantum_transport/calculations/localize.py
@@ -2,7 +2,6 @@
 
 import pickle
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import numpy as np
 from aiida import orm
@@ -11,17 +10,14 @@
 
 from .base import BaseCalculation
 
-if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
-
 
 class LocalizationCalculation(BaseCalculation):
     """docstring"""
 
     _default_parser_name = "quantum_transport.localize"
 
     @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
+    def define(cls, spec) -> None:
         """docstring"""
 
         super().define(spec)

diff --git a/src/aiida_quantum_transport/calculations/transmission.py b/src/aiida_quantum_transport/calculations/transmission.py
@@ -11,7 +11,7 @@
 from .base import BaseCalculation
 
 if TYPE_CHECKING:
-    from aiida.engine.processes.calcjobs.calcjob import CalcJobProcessSpec
+    from aiida.engine.processes.process_spec import CalcJobProcessSpec
 
 
 class TransmissionCalculation(BaseCalculation):