Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor systems #660

Merged
merged 12 commits into from
Oct 4, 2023
Prev Previous commit
Next Next commit
Avoid forward-refs where possible using TYPE_CHECKING
  • Loading branch information
stefsmeets committed Oct 3, 2023
commit 98d0452ef13b0d801691309aa227b705affb1e7c
6 changes: 4 additions & 2 deletions src/duqtools/_plot_utils.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
from __future__ import annotations

from typing import Union
from typing import TYPE_CHECKING, Union

import altair as alt
import numpy as np
import pandas as pd
import xarray as xr

if TYPE_CHECKING:
import pandas as pd


def _standardize_data(source: Union[pd.DataFrame, xr.Dataset],
z: str = 'time') -> pd.DataFrame:
Expand Down
6 changes: 5 additions & 1 deletion src/duqtools/cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,16 @@
import os
import shutil
from pathlib import Path
from typing import TYPE_CHECKING

from .config import Config
from .ids import ImasHandle
from .models import Locations, Run
from .models import Locations
from .operations import op_queue

if TYPE_CHECKING:
from .models import Run

logger = logging.getLogger(__name__)


Expand Down
8 changes: 5 additions & 3 deletions src/duqtools/ids/_apply_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,16 @@

import logging
from functools import partial
from typing import Union
from typing import TYPE_CHECKING, Union

import numpy as np

from .._logging_utils import duqlog_screen
from ..schema import IDSOperation
from ._handle import ImasHandle
from ._mapping import IDSMapping

if TYPE_CHECKING:
from ..schema import IDSOperation
from ._mapping import IDSMapping

logger = logging.getLogger(__name__)

Expand Down
4 changes: 3 additions & 1 deletion src/duqtools/ids/_handle.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

from ..config import lookup_vars, var_lookup
from ..operations import add_to_op_queue
from ..schema import IDSVariableModel, ImasBaseModel
from ..schema import ImasBaseModel
from ._copy import copy_ids_entry
from ._imas import imas, imasdef
from ._mapping import IDSMapping
Expand All @@ -21,6 +21,8 @@
if TYPE_CHECKING:
import xarray as xr

from ..schema import IDSVariableModel

logger = logging.getLogger(__name__)

_FILENAME = 'ids_{shot}{run:04d}{suffix}'
Expand Down
8 changes: 5 additions & 3 deletions src/duqtools/ids/_merge.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,19 @@
from __future__ import annotations

import logging
from typing import Sequence
from typing import TYPE_CHECKING, Sequence

import xarray as xr

from ..config import var_lookup
from ..operations import add_to_op_queue
from ..schema import IDSVariableModel
from ..utils import groupby
from ._handle import ImasHandle
from ._rebase import rebase_all_coords, squash_placeholders

if TYPE_CHECKING:
from ..schema import IDSVariableModel
from ._handle import ImasHandle

logger = logging.getLogger(__name__)

info = logger.info
Expand Down
7 changes: 4 additions & 3 deletions src/duqtools/ids/_rebase.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
from __future__ import annotations

import logging
from typing import Optional, Sequence, Union
from typing import TYPE_CHECKING, Optional, Sequence, Union

import numpy as np
import xarray as xr
if TYPE_CHECKING:
import numpy as np
import xarray as xr

logger = logging.getLogger(__name__)

Expand Down
3 changes: 1 addition & 2 deletions src/duqtools/jetto/_batchfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,8 @@

import jetto_tools

from ..config import Config

if TYPE_CHECKING:
from ..config import Config
from ..models import Job


Expand Down
6 changes: 5 additions & 1 deletion src/duqtools/jetto/_jettovar_to_json.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
from __future__ import annotations

import json
from typing import TYPE_CHECKING

from ..schema import JettoVar, JsetField, NamelistField
from ..schema import JsetField, NamelistField

if TYPE_CHECKING:
from ..schema import JettoVar


def jettovar_to_json(variable: JettoVar):
Expand Down
9 changes: 6 additions & 3 deletions src/duqtools/jetto/_system.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,22 +8,25 @@
import sys
from collections.abc import Sequence
from pathlib import Path
from typing import Any, List, Optional
from typing import TYPE_CHECKING, Any, List, Optional

from jetto_tools import config, jset, lookup, namelist, template
from jetto_tools import job as jetto_job
from jetto_tools.template import _EXTRA_FILE_REGEXES

from ..ids import ImasHandle
from ..jintrac import V210921Mixin, V220922Mixin
from ..models import AbstractSystem, Job
from ..models import AbstractSystem
from ..operations import add_to_op_queue
from ..schema import JettoVar
from ._batchfile import write_array_batchfile as _write_array_batchfile
from ._batchfile import write_batchfile as _write_batchfile
from ._jettovar_to_json import jettovar_to_json
from ._schema import JettoSystemModel

if TYPE_CHECKING:
from ..models import Job
from ..schema import JettoVar

if sys.version_info < (3, 10):
from importlib_resources import files
else:
Expand Down
6 changes: 5 additions & 1 deletion src/duqtools/list_variables.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,14 @@
from __future__ import annotations

import typing
from typing import TYPE_CHECKING

import click

from duqtools.config import Config, var_lookup
from duqtools.config import var_lookup

if TYPE_CHECKING:
from duqtools.config import Config

cs = click.style

Expand Down
6 changes: 4 additions & 2 deletions src/duqtools/merge.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,16 @@
from __future__ import annotations

import logging
from typing import Sequence
from typing import TYPE_CHECKING, Sequence

from .config import var_lookup
from .ids import ImasHandle, merge_data
from .operations import op_queue
from .schema import IDSVariableModel
from .utils import read_imas_handles_from_file

if TYPE_CHECKING:
from .schema import IDSVariableModel

logger = logging.getLogger(__name__)
info, debug = logger.info, logger.debug

Expand Down
4 changes: 3 additions & 1 deletion src/duqtools/models/_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@
import time
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING

import click

from ..config import Config
if TYPE_CHECKING:
from ..config import Config

logger = logging.getLogger(__name__)
info, debug = logger.info, logger.debug
Expand Down
9 changes: 6 additions & 3 deletions src/duqtools/models/_locations.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,15 @@
from __future__ import annotations

from pathlib import Path
from typing import Optional
from typing import TYPE_CHECKING, Optional

from pydantic_yaml import parse_yaml_raw_as

from ..config import Config
from ._run import Run, Runs
from ._run import Runs

if TYPE_CHECKING:
from ..config import Config
from ._run import Run


class Locations:
Expand Down
3 changes: 2 additions & 1 deletion src/duqtools/models/_system.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,13 @@
from typing import TYPE_CHECKING, Sequence

from ..config import Config
from ..schema import BaseModel, ImasBaseModel
from ..schema import BaseModel

if TYPE_CHECKING:
from pathlib import Path

from ..ids import ImasHandle
from ..schema import ImasBaseModel
from ._job import Job

logger = logging.getLogger(__name__)
Expand Down
5 changes: 4 additions & 1 deletion src/duqtools/sync_prominence.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,14 @@

import os
import subprocess
from typing import TYPE_CHECKING

from .config import Config
from .models import Job, Locations
from .operations import add_to_op_queue, op_queue

if TYPE_CHECKING:
from .config import Config


@add_to_op_queue('Getting data', 'job {job.path.name} from prominence')
def get_data_from_prominence(job: Job):
Expand Down
6 changes: 4 additions & 2 deletions src/duqtools/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,17 @@

import os
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any

from .config import Config
from .ets import Ets6System
from .jetto import JettoSystemV210921, JettoSystemV220922
from .jintrac import V220922Mixin
from .models import AbstractSystem
from .schema import NoSystemModel

if TYPE_CHECKING:
from .config import Config


class NoSystem(NoSystemModel, V220922Mixin, AbstractSystem):
"""This system is intended for workflows that need to apply some operations
Expand Down
3 changes: 1 addition & 2 deletions src/duqtools/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,8 @@

from pydantic_yaml import parse_yaml_raw_as

from ._types import PathLike

if TYPE_CHECKING:
from ._types import PathLike
from .ids import ImasHandle


Expand Down