Source code for autojob.study_group

"""Create study groups."""

import datetime
import json
from pathlib import Path
import shutil
from tempfile import TemporaryDirectory
from typing import Any
from typing import ClassVar
from typing import Self

from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from pydantic import FieldSerializationInfo
from pydantic import SerializerFunctionWrapHandler
from pydantic import field_serializer
from shortuuid import uuid

from autojob import SETTINGS
from autojob.study import Study
from autojob.utils.schemas import space_capitalize




class StudyGroup(BaseModel):
    """A collection of studies."""

    date_created: datetime.datetime = datetime.datetime.now(tz=datetime.UTC)
    study_group_id: str = Field(
        default_factory=lambda: "g" + uuid()[:9], alias="Study Group ID"
    )
    studies: list[Study] = Field(
        default=[],
    )
    name: str = ""
    notes: str = ""

    model_config: ClassVar = ConfigDict(
        populate_by_name=True, alias_generator=space_capitalize
    )



    @field_serializer(
        "studies", mode="wrap", return_type=list[Study] | list[str]
    )
    def serialize_tasks(
        self,
        v: Any,
        _: SerializerFunctionWrapHandler,
        info: FieldSerializationInfo,
    ) -> list[Study] | list[str]:
        """Serialize the studies in the study group."""
        if info.mode == "json":
            return [str(s.study_id) for s in self.studies]

        return v




    @classmethod
    def from_directory(
        cls,
        dir_name: Path,
        *,
        strict_mode: bool = SETTINGS.STRICT_MODE,
        legacy_mode: bool = False,
    ) -> Self:
        """Create a study group from a directory.

        Args:
            dir_name: The directory of a study group.
            strict_mode: Whether or not to require all outputs. If True,
                errors will be thrown on missing outputs. Defaults to
                ``SETTINGS.STRICT_MODE``.
            legacy_mode: Whether or not use the legacy mode directory structure.
                Defaults to False.
        """
        metadata_file = dir_name.joinpath(SETTINGS.STUDY_GROUP_FILE)
        with metadata_file.open(mode="r", encoding="utf-8") as file:
            metadata: dict[str, Any] = json.load(file)

        key = "Studies"

        studies: list[Study] = []

        for study in metadata[key]:
            source = dir_name.joinpath(study)
            studies.append(
                Study.from_directory(
                    source, strict_mode=strict_mode, legacy_mode=legacy_mode
                )
            )

        metadata[key] = studies

        return cls(**metadata)




    def to_directory(
        self, dir_name: Path, *, legacy_mode: bool = False
    ) -> None:
        """Create a directory for a study group.

        Args:
            dir_name: The directory in which to dump the :class:`StudyGroup`.
            legacy_mode: Whether or not use the legacy mode directory structure.
                Defaults to False.
        """
        with TemporaryDirectory() as tmpdir:
            metadata = self.model_dump(
                mode="json",
                by_alias=True,
            )

            study_group_path = Path(tmpdir).joinpath(self.study_group_id)
            study_group_path.mkdir()

            for study in sorted(self.studies, key=lambda s: s.study_id):
                dest = study_group_path.joinpath(study.study_id)
                study.to_directory(
                    dir_name=dest,
                    legacy_mode=legacy_mode,
                )

            with study_group_path.joinpath(SETTINGS.STUDY_GROUP_FILE).open(
                mode="w", encoding="utf-8"
            ) as file:
                json.dump(metadata, file, indent=4)

            shutil.copytree(study_group_path, dir_name, dirs_exist_ok=True)