Source code for ewoksfluo.tasks.input.pick_scans

from typing import List
from typing import Optional
from typing import Tuple

try:
    from typing import Self
except ImportError:
    from typing_extensions import Self

from ewokscore import Task
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from pydantic import Field
from pydantic import model_validator

from .pick_utils import pick_scans




class Inputs(BaseInputModel):
    filenames: List[str] = Field(
        description="Bliss dataset HDF5 file names.",
        examples=[
            ["/data/dataset1.h5", "/data/dataset2.h5"],
        ],
    )

    scan_ranges: List[Tuple[int, int]] = Field(
        description="Ranges of scan numbers.",
        examples=[
            [(1, 4), (100, 101)],
        ],
    )

    exclude_scans: List[List[int]] = Field(
        default_factory=list,
        description="Scan numbers to exclude for each range.",
        examples=[[[1, 3], [1]]],
    )

    retry_timeout: Optional[float] = Field(
        20,
        description="Timeout in seconds when waiting for the scan to be fully written. "
        "`None` means wait forever. Negative means do not wait.",
    )
    retry_period: float = Field(
        0.5,
        description="Retry period in seconds when waiting for the scan to be fully written.",
    )



    @model_validator(mode="after")
    def check_and_expand_lengths(self) -> Self:
        n_files = len(self.filenames)
        n_ranges = len(self.scan_ranges)
        n_exclude = len(self.exclude_scans)

        if n_files != n_ranges:
            raise ValueError("`filenames` and `scan_ranges` must have the same length")

        if n_exclude == 0:
            self.exclude_scans = [[] for _ in range(n_files)]
            return self

        if n_exclude != n_files:
            raise ValueError(
                "`exclude_scans` must have the same length as `filenames` or be empty"
            )

        return self






class Outputs(BaseOutputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[
            [
                "/data/dataset1.h5::/2.1",
                "/data/dataset1.h5::/4.1",
                "/data/dataset2.h5::/100.1",
                "/data/dataset2.h5::/101.1",
            ]
        ],
    )





class PickScans(Task, input_model=Inputs, output_model=Outputs):
    """Select a multiple Bliss scan from multiple files."""



    def run(self):
        self.outputs.bliss_scan_uris = pick_scans(
            self.inputs.filenames, self.inputs.scan_ranges, self.inputs.exclude_scans
        )