Skip to content

Commit afed5dc

Browse files
authoredJun 17, 2023
🔧 Add types for DefaultFieldsAttributeDict subclasses (#6059)
Improves type checking and LSP completion.
1 parent 3976344 commit afed5dc

File tree

3 files changed

+102
-2
lines changed

3 files changed

+102
-2
lines changed
 

‎aiida/common/datastructures.py

+38
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88
# For further information please visit http://www.aiida.net #
99
###########################################################################
1010
"""Module to define commonly used data structures."""
11+
from __future__ import annotations
12+
1113
from enum import Enum, IntEnum
14+
from typing import TYPE_CHECKING
1215

1316
from .extendeddicts import DefaultFieldsAttributeDict
1417

@@ -93,6 +96,31 @@ class CalcInfo(DefaultFieldsAttributeDict):
9396
'provenance_exclude_list', 'codes_info', 'codes_run_mode', 'skip_submit'
9497
)
9598

99+
if TYPE_CHECKING:
100+
101+
job_environment: None | dict[str, str]
102+
email: None | str
103+
email_on_started: bool
104+
email_on_terminated: bool
105+
uuid: None | str
106+
prepend_text: None | str
107+
append_text: None | str
108+
num_machines: None | int
109+
num_mpiprocs_per_machine: None | int
110+
priority: None | int
111+
max_wallclock_seconds: None | int
112+
max_memory_kb: None | int
113+
rerunnable: bool
114+
retrieve_list: None | list[str | tuple[str, str, str]]
115+
retrieve_temporary_list: None | list[str | tuple[str, str, str]]
116+
local_copy_list: None | list[tuple[str, str, str]]
117+
remote_copy_list: None | list[tuple[str, str, str]]
118+
remote_symlink_list: None | list[tuple[str, str, str]]
119+
provenance_exclude_list: None | list[str]
120+
codes_info: None | list[CodeInfo]
121+
codes_run_mode: None | CodeRunMode
122+
skip_submit: None | bool
123+
96124

97125
class CodeInfo(DefaultFieldsAttributeDict):
98126
"""
@@ -148,6 +176,16 @@ class CodeInfo(DefaultFieldsAttributeDict):
148176
'code_uuid'
149177
)
150178

179+
if TYPE_CHECKING:
180+
181+
cmdline_params: None | list[str]
182+
stdin_name: None | str
183+
stdout_name: None | str
184+
stderr_name: None | str
185+
join_files: None | bool
186+
withmpi: None | bool
187+
code_uuid: None | str
188+
151189

152190
class CodeRunMode(IntEnum):
153191
"""Enum to indicate the way the codes of a calculation should be run.

‎aiida/engine/processes/calcjobs/calcjob.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -964,7 +964,7 @@ def presubmit(self, folder: Folder) -> CalcInfo:
964964
tmpl_code_info.stdin_name = code_info.stdin_name
965965
tmpl_code_info.stdout_name = code_info.stdout_name
966966
tmpl_code_info.stderr_name = code_info.stderr_name
967-
tmpl_code_info.join_files = code_info.join_files
967+
tmpl_code_info.join_files = code_info.join_files or False
968968

969969
tmpl_codes_info.append(tmpl_code_info)
970970
job_tmpl.codes_info = tmpl_codes_info

‎aiida/schedulers/datastructures.py

+63-1
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@
2121
from datetime import datetime, timezone
2222
import enum
2323
import json
24+
from typing import TYPE_CHECKING
2425

25-
from aiida.common import AIIDA_LOGGER
26+
from aiida.common import AIIDA_LOGGER, CodeRunMode
2627
from aiida.common.extendeddicts import AttributeDict, DefaultFieldsAttributeDict
2728
from aiida.common.timezone import make_aware, timezone_from_name
2829

@@ -108,6 +109,12 @@ class NodeNumberJobResource(JobResource):
108109
'num_cores_per_mpiproc',
109110
)
110111

112+
if TYPE_CHECKING:
113+
num_machines: int
114+
num_mpiprocs_per_machine: int
115+
num_cores_per_machine: int
116+
num_cores_per_mpiproc: int
117+
111118
@classmethod
112119
def validate_resources(cls, **kwargs):
113120
"""Validate the resources against the job resource class of this scheduler.
@@ -193,6 +200,10 @@ class ParEnvJobResource(JobResource):
193200
'tot_num_mpiprocs',
194201
)
195202

203+
if TYPE_CHECKING:
204+
parallel_env: str
205+
tot_num_mpiprocs: int
206+
196207
@classmethod
197208
def validate_resources(cls, **kwargs):
198209
"""Validate the resources against the job resource class of this scheduler.
@@ -366,6 +377,34 @@ class JobTemplate(DefaultFieldsAttributeDict): # pylint: disable=too-many-insta
366377
'codes_info',
367378
)
368379

380+
if TYPE_CHECKING:
381+
shebang: str
382+
submit_as_hold: bool
383+
rerunnable: bool
384+
job_environment: dict[str, str] | None
385+
environment_variables_double_quotes: bool | None
386+
working_directory: str
387+
email: str
388+
email_on_started: bool
389+
email_on_terminated: bool
390+
job_name: str
391+
sched_output_path: str
392+
sched_error_path: str
393+
sched_join_files: bool
394+
queue_name: str
395+
account: str
396+
qos: str
397+
job_resource: JobResource
398+
priority: str
399+
max_memory_kb: int | None
400+
max_wallclock_seconds: int
401+
custom_scheduler_commands: str
402+
prepend_text: str
403+
append_text: str
404+
import_sys_environment: bool | None
405+
codes_run_mode: CodeRunMode
406+
codes_info: list[JobTemplateCodeInfo]
407+
369408

370409
@dataclass
371410
class JobTemplateCodeInfo:
@@ -474,6 +513,29 @@ class JobInfo(DefaultFieldsAttributeDict): # pylint: disable=too-many-instance-
474513
'finish_time'
475514
)
476515

516+
if TYPE_CHECKING:
517+
job_id: str
518+
title: str
519+
exit_status: int
520+
terminating_signal: int
521+
annotation: str
522+
job_state: JobState
523+
job_substate: str
524+
allocated_machines: list[MachineInfo]
525+
job_owner: str
526+
num_mpiprocs: int
527+
num_cpus: int
528+
num_machines: int
529+
queue_name: str
530+
account: str
531+
qos: str
532+
wallclock_time_seconds: int
533+
requested_wallclock_time_seconds: int
534+
cpu_time: int
535+
submission_time: datetime
536+
dispatch_time: datetime
537+
finish_time: datetime
538+
477539
# If some fields require special serializers, specify them here.
478540
# You then need to define also the respective _serialize_FIELDTYPE and
479541
# _deserialize_FIELDTYPE methods

0 commit comments

Comments
 (0)
Please sign in to comment.