mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-11 14:49:43 +00:00
63dca86df8
* fix(root): adding proper support for the root and new_root arguments * feat(roots): adding a roots agrument for the merge operation * chore(clean): cleaning up code * chore(doctrings): updating doctrings with new features * fix(repo_id): setting repo_id to None when not needed * fix(roots/repo_ids): making mypy happy by using repo_ids and roots for merge operation * fix(path): fixing path related issues * fix(repo_id): fixing issues related to repo_id * chore(doctrings): updating docstrings + fix typo * chore(clean): cleaning code * fix(split new_repo_id): reverting new_repo_id addition for split operation * docs(dosctrings): completing docstrings * fix(repo_ids/roots): improving checks for repo_ids/roots lengths * fix(repo_ids): making repo_ids optional in MergeConfig but raise if not given * fix(docstrings): fixing docstrings for split operation * fix(hints): updating get_output_path hints to accept paths as strings too * fix(y/N prompts): removing y/N prompts in lerobot_edit_dataset * fix(merge repo_id): fixing merge operation to use new_repo_id instead of repo_id * fix(typo): fixing typo in doctrings
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import draccus
|
|
import pytest
|
|
|
|
from lerobot.scripts.lerobot_edit_dataset import (
|
|
ConvertImageToVideoConfig,
|
|
DeleteEpisodesConfig,
|
|
EditDatasetConfig,
|
|
InfoConfig,
|
|
MergeConfig,
|
|
ModifyTasksConfig,
|
|
OperationConfig,
|
|
RemoveFeatureConfig,
|
|
SplitConfig,
|
|
_validate_config,
|
|
)
|
|
|
|
|
|
def parse_cfg(cli_args: list[str]) -> EditDatasetConfig:
|
|
"""Helper to parse CLI args into an EditDatasetConfig via draccus."""
|
|
return draccus.parse(EditDatasetConfig, args=cli_args)
|
|
|
|
|
|
class TestOperationTypeParsing:
|
|
"""Test that --operation.type correctly selects the right config subclass."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"type_name, expected_cls",
|
|
[
|
|
("delete_episodes", DeleteEpisodesConfig),
|
|
("split", SplitConfig),
|
|
("merge", MergeConfig),
|
|
("remove_feature", RemoveFeatureConfig),
|
|
("modify_tasks", ModifyTasksConfig),
|
|
("convert_image_to_video", ConvertImageToVideoConfig),
|
|
("info", InfoConfig),
|
|
],
|
|
)
|
|
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
|
cfg = parse_cfg(
|
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
|
)
|
|
assert isinstance(cfg.operation, expected_cls), (
|
|
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
|
)
|
|
|
|
def test_merge_requires_new_repo_id(self):
|
|
cfg = parse_cfg(["--operation.type", "merge"])
|
|
with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
|
|
_validate_config(cfg)
|
|
|
|
def test_non_merge_requires_repo_id(self):
|
|
cfg = parse_cfg(["--operation.type", "delete_episodes"])
|
|
with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):
|
|
_validate_config(cfg)
|
|
|
|
@pytest.mark.parametrize(
|
|
"type_name, expected_cls",
|
|
[
|
|
("delete_episodes", DeleteEpisodesConfig),
|
|
("split", SplitConfig),
|
|
("merge", MergeConfig),
|
|
("remove_feature", RemoveFeatureConfig),
|
|
("modify_tasks", ModifyTasksConfig),
|
|
("convert_image_to_video", ConvertImageToVideoConfig),
|
|
("info", InfoConfig),
|
|
],
|
|
)
|
|
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
|
cfg = parse_cfg(
|
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
|
)
|
|
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
|
assert resolved_name == type_name
|