Skip to content

Commit

Permalink
Use Python 3.9 compatible type hints for pydantic objects
Browse files Browse the repository at this point in the history
  • Loading branch information
ml-evs committed Oct 11, 2023
1 parent da0cb7a commit 8853680
Showing 1 changed file with 22 additions and 22 deletions.
44 changes: 22 additions & 22 deletions src/jobflow_remote/config/base.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
from __future__ import annotations

import abc
import logging
import traceback
from enum import Enum
from pathlib import Path
from typing import Annotated, Literal, Union
from typing import Annotated, Literal, Optional, Union

from jobflow import JobStore
from pydantic import BaseModel, Extra, Field, validator
Expand Down Expand Up @@ -36,7 +34,7 @@ class RunnerOptions(BaseModel):
30,
description="Delay between subsequent advancement of the job's remote state (seconds)",
)
lock_timeout: int | None = Field(
lock_timeout: Optional[int] = Field(
86400,
description="Time to consider the lock on a document expired and can be overridden (seconds)",
)
Expand Down Expand Up @@ -115,16 +113,16 @@ class WorkerBase(BaseModel):
description="Absolute path of the directory of the worker where subfolders for "
"executing the calculation will be created"
)
resources: dict | None = Field(
resources: Optional[dict] = Field(
None,
description="A dictionary defining the default resources requested to the "
"scheduler. Used to fill in the QToolKit template",
)
pre_run: str | None = Field(
pre_run: Optional[str] = Field(
None,
description="String with commands that will be executed before the execution of the Job",
)
post_run: str | None = Field(
post_run: Optional[str] = Field(
None,
description="String with commands that will be executed after the execution of the Job",
)
Expand Down Expand Up @@ -234,13 +232,13 @@ class RemoteWorker(WorkerBase):
host: str = Field(description="The host to which to connect")
user: str = Field(None, description="Login username")
port: int = Field(None, description="Port number")
password: str | None = Field(None, description="Login password")
key_filename: str | list[str] | None = Field(
password: Optional[str] = Field(None, description="Login password")
key_filename: Optional[Union[str, list[str]]] = Field(
None,
description="The filename, or list of filenames, of optional private key(s) "
"and/or certs to try for authentication",
)
passphrase: str | None = Field(
passphrase: Optional[str] = Field(
None, description="Passphrase used for decrypting private keys"
)
gateway: str = Field(
Expand All @@ -259,10 +257,10 @@ class RemoteWorker(WorkerBase):
description="Whether to send environment variables 'inline' as prefixes in "
"front of command strings",
)
keepalive: int | None = Field(
keepalive: Optional[int] = Field(
60, description="Keepalive value in seconds passed to paramiko's transport"
)
shell_cmd: str | None = Field(
shell_cmd: Optional[str] = Field(
"bash",
description="The shell command used to execute the command remotely. If None "
"the command is executed directly",
Expand Down Expand Up @@ -326,14 +324,16 @@ class ExecutionConfig(BaseModel):
Configuration to be set before and after the execution of a Job.
"""

modules: list[str] | None = Field(None, description="list of modules to be loaded")
export: dict[str, str] | None = Field(
modules: Optional[list[str]] = Field(
None, description="list of modules to be loaded"
)
export: Optional[dict[str, str]] = Field(
None, description="dictionary with variable to be exported"
)
pre_run: str | None = Field(
pre_run: Optional[str] = Field(
None, description="Other commands to be executed before the execution of a job"
)
post_run: str | None = Field(
post_run: Optional[str] = Field(
None, description="Commands to be executed after the execution of a job"
)

Expand All @@ -347,20 +347,20 @@ class Project(BaseModel):
"""

name: str = Field(description="The name of the project")
base_dir: str | None = Field(
base_dir: Optional[str] = Field(
None,
description="The base directory containing the project related files. Default "
"is a folder with the project name inside the projects folder",
)
tmp_dir: str | None = Field(
tmp_dir: Optional[str] = Field(
None,
description="Folder where remote files are copied. Default a 'tmp' folder in base_dir",
)
log_dir: str | None = Field(
log_dir: Optional[str] = Field(
None,
description="Folder containing all the logs. Default a 'log' folder in base_dir",
)
daemon_dir: str | None = Field(
daemon_dir: Optional[str] = Field(
None,
description="Folder containing daemon related files. Default to a 'daemon' "
"folder in base_dir",
Expand Down Expand Up @@ -390,11 +390,11 @@ class Project(BaseModel):
description="The JobStore used for the input. Can contain the monty "
"serialized dictionary or the Store int the Jobflow format",
)
metadata: dict | None = Field(
metadata: Optional[dict] = Field(
None, description="A dictionary with metadata associated to the project"
)

def get_jobstore(self) -> JobStore | None:
def get_jobstore(self) -> Optional[JobStore]:
"""
Generate an instance of the JobStore based on the configuration
Expand Down

0 comments on commit 8853680

Please sign in to comment.