Source code for autojob.study_group

"""Create study groups."""

from datetime import UTC
from datetime 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 UUID4
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 autojob import SETTINGS
from autojob.study import Study
from autojob.utils.schemas import id_factory

STUDY_GROUP_FIELDS = [
    "label",
    "study_group_id",
    "date_created",
]




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

    date_created: datetime = Field(
        default_factory=lambda: datetime.now(tz=UTC),
        description="The date and time that the study group was created",
    )
    study_group_id: UUID4 | str = Field(
        default_factory=id_factory("g"),
        description="A unique identifier for the study group",
        union_mode="left_to_right",
    )
    studies: list[Study] = Field(
        default=[], description="The studies in the study group"
    )
    name: str = Field(default="", description="The name of the study group")
    notes: list[str] = Field(
        default=[], description="A list of notes about the study group"
    )

    model_config: ClassVar = ConfigDict(populate_by_name=True)



    @field_serializer(
        "studies", mode="wrap", return_type=list[Study] | list[str]
    )
    def serialize_studies(
        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 | None = None
    ) -> 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``.
        """
        strict_mode = (
            SETTINGS.STRICT_MODE if strict_mode is None else strict_mode
        )
        metadata_file = Path(dir_name, SETTINGS.STUDY_GROUP_METADATA_FILE)
        with metadata_file.open(mode="r", encoding="utf-8") as file:
            metadata: dict[str, Any] = json.load(file)

        studies: list[Study] = []

        for study in metadata["studies"]:
            source = Path(dir_name, study)
            studies.append(
                Study.from_directory(source, strict_mode=strict_mode)
            )

        metadata["studies"] = studies
        return cls(**metadata)




    def to_directory(self, src: Path) -> None:
        """Create a directory for a study group.

        Args:
            src: The directory in which to dump the :class:`StudyGroup`.
        """
        with TemporaryDirectory() as tmpdir:
            metadata = self.model_dump(mode="json")
            study_group_path = Path(tmpdir, str(self.study_group_id))
            study_group_path.mkdir()

            for study in sorted(self.studies, key=lambda s: s.study_id):
                dest = Path(study_group_path, str(study.study_id))
                study.to_directory(dest=dest)

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

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