Coverage for collectra/partition.py: 97.45%
157 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1from collections.abc import Iterator, Sequence
2from pathlib import Path
4from rich.progress import Progress
6from .commons.files import CollectraFile
9def partition_results(inputs, ext, partitions, seed=42, force=False):
10 """Resolve pipeline results, then use the existing partitioning engine.
12 Raises:
13 ValueError: If any result already has a partition and ``force`` is False.
14 """
15 from collectra.utils import resolve_files
17 if not partitions:
18 raise ValueError("Specify partitions, for example --train 80% --validation 20%")
19 files = []
20 for value in inputs:
21 path = Path(value).expanduser().resolve()
22 if not path.exists():
23 raise FileNotFoundError(f"Input does not exist: {path}")
24 if path.is_file() and path.name == "results.yaml":
25 files.append(path.parent)
26 else:
27 files.extend(resolve_files([path], [f".{ext.lower()}"]))
28 files = sorted(set(files))
29 if not files:
30 raise ValueError(f"No .{ext} results found in the supplied inputs")
31 # A single Progress renders its tasks as one table, so the bars line up.
32 with Progress() as progress:
33 # Validate all results before any partition metadata is changed.
34 already_partitioned = [
35 path
36 for path in _track(progress, files, "Checking results")
37 if CollectraFile.from_data(path).collectra_results_metadata.partition
38 is not None
39 ]
40 if already_partitioned and not force:
41 raise ValueError(
42 f"{len(already_partitioned)} of {len(files)} results already have a partition. "
43 "Use --force to override the current partitions."
44 )
45 process_partitions(files, partitions, seed, progress=progress)
48def _track(
49 progress: Progress | None, items: Sequence[Path], description: str
50) -> Iterator[Path]:
51 """Yields items, advancing a task on the shared progress display when one is given."""
52 if progress is None:
53 yield from items
54 return
55 task = progress.add_task(description, total=len(items))
56 for item in items:
57 yield item
58 progress.advance(task)
61def get_files(args: list[str]) -> tuple[list[Path], list[str]]:
62 """Gets the list of files from the command-line arguments. It also returns a reduced list of arguments that excludes the file paths, which can be used for further processing,
64 Args:
65 args (list[str]): The list of command-line arguments.
67 Raises:
68 ValueError: If an unexpected option is encountered.
69 ValueError: If input files have mismatched extensions.
70 Returns:
71 tuple[list[Path], list[str]]: A tuple containing the list of file paths and the reduced list of arguments.
72 """
74 reduced_args: list[str] = []
76 files: list[Path] = []
77 idx_skip = -1
78 for idx, arg in enumerate(args):
79 if idx == idx_skip:
80 reduced_args.append(arg)
81 continue
82 if arg.startswith("--"):
83 idx_skip = idx + 1
84 reduced_args.append(arg)
85 continue
86 if arg.startswith("-"):
87 raise ValueError(f"Unexpected option {arg}")
88 path = Path(arg)
89 if path.exists():
90 files.append(path)
92 current_suffix = ""
93 for file in files:
94 if current_suffix == "":
95 current_suffix = file.suffix
96 elif file.suffix != current_suffix:
97 raise ValueError(
98 f"All input files must have the same extension. Found mismatch at {file.name}"
99 )
101 return files, reduced_args
104def get_partitions(args: list[str]) -> dict[str, str | int | float]:
105 """Extracts partition specifications from the command-line arguments.
107 Args:
108 args (list[str]): The list of command-line arguments.
110 Raises:
111 ValueError: If a partition value is missing or invalid.
113 Returns:
114 dict[str, str | int | float]: A dictionary mapping partition labels to their values.
115 """
116 partitions: dict[str, str | int | float] = {}
117 for idx, arg in enumerate(args):
118 if arg.startswith("--"):
119 label = arg.replace("--", "")
120 idx += 1 # Move to the next argument which should be the value for this partition
121 value = args[idx] if idx < len(args) else None
122 if not value or value.startswith("-"):
123 raise ValueError(
124 f"Expected a value for partition {label}, but got {value}"
125 )
126 partitions[label] = value
127 return partitions
130def get_partition_value_type(value: str | int | float) -> str:
131 """Determines the type of the partition value (percentage, integer, or float) based on its format.
133 Args:
134 value (str | int | float): The partition value to be evaluated.
136 Raises:
137 ValueError: If the value does not match any expected type.
138 """
139 if isinstance(value, str) and value.endswith("%"):
140 return "percentage"
141 try:
142 int(value) # Check if it can be converted to an integer
143 return "integer"
144 except ValueError:
145 try:
146 float(value) # Check if it can be converted to a float
147 return "float"
148 except ValueError:
149 raise ValueError(
150 f"Invalid value for partition: {value}. Must be an integer, float, or percentage."
151 )
154def validate_partition_value(
155 value: str | int | float, expected_type: str
156) -> int | float:
157 """Validates and converts the partition value to the appropriate type based on the expected type.
159 Args:
160 value (str | int | float): The partition value to be validated.
161 expected_type (str): The expected type of the partition value ("percentage", "integer", or "float").
163 Raises:
164 ValueError: If the value does not match the expected type or format.
166 Returns:
167 int | float: The validated and converted partition value.
168 """
169 match expected_type:
170 case "percentage":
171 if not (isinstance(value, str) and value.endswith("%")):
172 raise ValueError(f"Expected a percentage value, but got {value}")
173 return float(
174 int(value.replace("%", "")) / 100
175 ) # Remove the percentage sign for further processing
176 case "integer":
177 return int(
178 value
179 ) # Convert to integer if it's a string representation of an integer
180 case "float":
181 return float(
182 value
183 ) # Convert to float if it's a string representation of a float
184 case _:
185 raise ValueError(f"Unknown expected type: {expected_type}")
188def validate_partition_sum(
189 validated_partitions: dict[str, int | float],
190 value_type: str | None,
191 files: list[Path],
192):
193 """Validates that the sum of partition values is appropriate based on the value type (percentage, integer, or float).
195 Args:
196 validated_partitions (dict[str, int | float]): The dictionary of validated partition values.
197 value_type (str | None): The type of the partition values ("percentage", "integer", or "float").
198 files (list[Path]): The list of files to be partitioned
200 Raises:
201 ValueError: If the total percentage exceeds 100%.
202 ValueError: If the total count of partitions exceeds the number of files.
203 ValueError: If the total fraction of partitions exceeds 1.
204 ValueError: If the value type is unknown.
206 """
207 values: list[int | float] = list(validated_partitions.values())
208 match value_type:
209 case "percentage":
210 total_percentage = sum(values)
211 if total_percentage > 1.0:
212 raise ValueError(
213 f"Total percentage of partitions must not exceed 100%. Currently, it sums to {total_percentage * 100}%."
214 )
215 case "integer":
216 total_count = sum(values)
217 if total_count > len(files):
218 raise ValueError(
219 f"Total count of partitions ({total_count}) cannot exceed the number of files ({len(files)})."
220 )
221 case "float":
222 total_fraction = sum(values)
223 if total_fraction > 1.0:
224 raise ValueError(
225 f"Total fraction of partitions cannot exceed 1. Currently, it sums to {total_fraction}."
226 )
227 case _:
228 raise ValueError(f"Unknown value type: {value_type}")
231def assign_partition_label(file: Path, label: str | None):
232 """Assigns the specified partition label to the given file by updating its metadata.
234 Args:
235 file (Path): The file to which the partition label will be assigned.
236 label (str): The partition label to be assigned to the file.
238 """
239 collectra_file = CollectraFile.from_data(file)
240 collectra_file.collectra_results_metadata.partition = label
241 collectra_file.save()
244def assign_partition_label_to_files(
245 files: list[Path], label: str, progress: Progress | None = None
246):
247 """Assigns the specified partition label to the given list of files by updating their metadata.
249 Args:
250 files (list[Path]): The list of files to which the partition label will be assigned.
251 label (str): The partition label to be assigned to the files.
252 progress (Progress | None): Shared progress display to report on, if any.
254 Raises:
255 Exception: If there is an error in assigning the partition label to any of the files, reverts any changes made to the files before the error occurred to maintain consistency.
257 """
258 old_labels: list[tuple[Path, str | None]] = []
259 try:
260 for file in _track(
261 progress, files, f"Partition '{label}' ({len(files)} files)"
262 ):
263 collectra_file = CollectraFile.from_data(file)
264 metadata = collectra_file.collectra_results_metadata
265 old_labels.append((file, metadata.partition))
266 metadata.partition = label
267 collectra_file.save()
268 except Exception as e:
269 # Revert any changes made to files in case of an error
270 for file, old_label in old_labels:
271 assign_partition_label(file, old_label)
272 raise e
275def partition_files(
276 files: list[Path],
277 partitions: dict[str, int | float],
278 value_type: str | None,
279 seed: int,
280 progress: Progress | None = None,
281):
282 """Partitions the list of files based on the specified partition values and their type. It shuffles the files using the provided random seed to ensure reproducibility.
284 Args:
285 files (list[Path]): The list of files to be partitioned.
286 partitions (dict[str, int | float]): The dictionary of partition labels and their corresponding values.
287 value_type (str | None): The type of the partition values ("percentage", "integer", or "float").
288 seed (int): The random seed for shuffling the files.
289 progress (Progress | None): Shared progress display to report on, if any.
291 Raises:
292 ValueError: If the value type is not determined before partitioning.
293 ValueError: If an unknown value type is encountered during partitioning.
295 """
296 import random
298 if value_type is None:
299 raise ValueError("Value type must be determined before partitioning files.")
301 random.seed(seed)
302 random.shuffle(files)
303 total_files = len(files)
305 for label, value in partitions.items():
306 match value_type:
307 case "percentage":
308 num_files = int(value * total_files)
309 case "integer":
310 num_files = int(value)
311 case "float":
312 num_files = int(value * total_files)
313 case _:
314 raise ValueError(f"Unknown value type: {value_type}")
315 partition_files = files[:num_files]
316 files = files[num_files:]
317 assign_partition_label_to_files(partition_files, label, progress=progress)
320def process_partitions(
321 files: list[Path],
322 partitions: dict[str, str | int | float],
323 seed: int,
324 progress: Progress | None = None,
325):
326 """Processes the partition specifications by validating their values, ensuring they sum up appropriately, and then partitioning the files accordingly.
328 Args:
329 files (list[Path]): The list of files to be partitioned.
330 partitions (dict[str, str | int | float]): The dictionary of partition labels and their corresponding values (as strings, integers, or floats).
331 seed (int): The random seed for shuffling the files.
332 progress (Progress | None): Shared progress display to report on, if any.
334 Raises:
335 ValueError: If there is an error in processing any of the partition specifications.
337 """
338 value_type: str | None = None
339 validated_partitions: dict[str, int | float] = {}
340 for label, value in partitions.items():
341 try:
342 if value_type is None:
343 value_type = get_partition_value_type(value)
344 validated_partitions[label] = validate_partition_value(value, value_type)
345 except ValueError as e:
346 raise ValueError(f"Error processing partition {label}: {e}")
347 validate_partition_sum(validated_partitions, value_type, files)
348 partition_files(files, validated_partitions, value_type, seed, progress=progress)