"""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.files import create_templated_dir_name
from autojob.utils.files import find_study_dirs
from autojob.utils.schemas import id_factory
STUDY_GROUP_FIELDS = [
"label",
"study_group_id",
"date_created",
]
[docs]
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)
[docs]
@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
[docs]
@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 find_study_dirs(dir_name):
source = Path(dir_name, study)
studies.append(
Study.from_directory(source, strict_mode=strict_mode)
)
metadata["studies"] = studies
return cls(**metadata)
[docs]
def to_directory(
self,
src: Path,
*,
study_group_template: str | None = None,
study_template: str | None = None,
task_group_template: str | None = None,
task_template: str | None = None,
) -> Path:
"""Create a directory for a study group.
Args:
src: The directory in which to dump the :class:`StudyGroup`.
study_group_template: A template string for naming study group
directories. Defaults to None in which case the study group
ID will be used to create the directory.
study_template: A template string for naming study directories.
Defaults to None in which case the study ID will be used to
create the directory.
task_group_template: A template string for naming task group
directories. Defaults to None in which case the task group ID
will be used to create the directory.
task_template: A template string for naming task directories.
Defaults to None in which case the task ID will be used to
create the directory.
Returns:
The study group directory that was created.
"""
with TemporaryDirectory() as tmpdir:
metadata = self.model_dump(mode="json")
if study_group_template:
dir_name = create_templated_dir_name(
study_group_template, Path(tmpdir), metadata
)
else:
dir_name = str(self.study_group_id)
study_group_path = Path(tmpdir, dir_name)
study_group_path.mkdir()
for study in sorted(self.studies, key=lambda s: s.study_id):
study.to_directory(
dest=study_group_path,
study_template=study_template,
task_group_template=task_group_template,
task_template=task_template,
)
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)
created = Path(src, dir_name)
shutil.copytree(
study_group_path,
created,
dirs_exist_ok=True,
)
return created