Coverage for src / autoencodix / data / _datasetcontainer.py: 94%

1from typing import Optional, Union 

2from dataclasses import dataclass 

3from autoencodix.base._base_dataset import BaseDataset 

4 

5 

6@dataclass 

7class DatasetContainer: 

8 """A container for datasets used in training, validation, and testing. 

9 

10 Attributes: 

11 train : Dataset 

12 The training dataset. 

13 valid : Dataset 

14 The validation dataset. 

15 test : Dataset 

16 The testing dataset. 

17 """ 

18 

19 train: Optional[Union[BaseDataset, None]] = None 

20 valid: Optional[Union[BaseDataset, None]] = None 

21 test: Optional[Union[BaseDataset, None]] = None 

22 

23 def __getitem__(self, key: str) -> BaseDataset: 

24 """Allows dictionary-like access to datasets.""" 

25 if key not in {"train", "valid", "test"}: 

26 raise KeyError(f"Invalid key: {key}. Must be 'train', 'valid', or 'test'.") 

27 return getattr(self, key) 

28 

29 def __setitem__(self, key: str, value: BaseDataset): 

30 """Allows dictionary-like assignment of datasets.""" 

31 if key not in {"train", "valid", "test"}: 

32 raise KeyError(f"Invalid key: {key}. Must be 'train', 'valid', or 'test'.") 

33 setattr(self, key, value)