feat(pipeline): Add __repr__ method to RobotProcessor for improved readability

- Implemented a __repr__ method in the RobotProcessor class to provide a clear string representation of the processor, including step names and optional parameters like name and seed.
- Added comprehensive tests to validate the __repr__ output for various scenarios, including empty processors, single and multiple steps, custom names, and seed values.
- Ensured that the representation handles long lists of steps with truncation for better readability.
This commit is contained in:
Adil Zouitine
2025-07-22 10:49:21 +02:00
parent 26cb9a24c3
commit 9fe3a3fb17
2 changed files with 144 additions and 0 deletions
+20
View File
@@ -729,6 +729,26 @@ class RobotProcessor(ModelHubMixin):
return profile_results
def __repr__(self) -> str:
"""Return a readable string representation of the processor."""
step_names = [step.__class__.__name__ for step in self.steps]
if not step_names:
steps_repr = "steps=0: []"
elif len(step_names) <= 3:
steps_repr = f"steps={len(step_names)}: [{', '.join(step_names)}]"
else:
# Show first 2 and last 1 with ellipsis for long lists
displayed = f"{step_names[0]}, {step_names[1]}, ..., {step_names[-1]}"
steps_repr = f"steps={len(step_names)}: [{displayed}]"
parts = [f"name='{self.name}'", steps_repr]
if self.seed is not None:
parts.append(f"seed={self.seed}")
return f"RobotProcessor({', '.join(parts)})"
class ObservationProcessor:
"""Base class for processors that modify only the observation component of a transition.
+124
View File
@@ -1531,3 +1531,127 @@ def test_to_device_with_mixed_state_types():
# Move back to CPU
pipeline.to("cpu")
assert step.tensor_data.device.type == "cpu"
def test_repr_empty_processor():
"""Test __repr__ with empty processor."""
pipeline = RobotProcessor()
repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=0: [])"
assert repr_str == expected
def test_repr_single_step():
"""Test __repr__ with single step."""
step = MockStep("test_step")
pipeline = RobotProcessor([step])
repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep])"
assert repr_str == expected
def test_repr_multiple_steps_under_limit():
"""Test __repr__ with 2-3 steps (all shown)."""
step1 = MockStep("step1")
step2 = MockStepWithoutOptionalMethods()
pipeline = RobotProcessor([step1, step2])
repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])"
assert repr_str == expected
# Test with 3 steps (boundary case)
step3 = MockStepWithTensorState()
pipeline = RobotProcessor([step1, step2, step3])
repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=3: [MockStep, MockStepWithoutOptionalMethods, MockStepWithTensorState])"
assert repr_str == expected
def test_repr_many_steps_truncated():
"""Test __repr__ with more than 3 steps (truncated with ellipsis)."""
step1 = MockStep("step1")
step2 = MockStepWithoutOptionalMethods()
step3 = MockStepWithTensorState()
step4 = MockModuleStep()
step5 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4, step5])
repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=5: [MockStep, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])"
assert repr_str == expected
def test_repr_with_custom_name():
"""Test __repr__ with custom processor name."""
step = MockStep("test_step")
pipeline = RobotProcessor([step], name="CustomProcessor")
repr_str = repr(pipeline)
expected = "RobotProcessor(name='CustomProcessor', steps=1: [MockStep])"
assert repr_str == expected
def test_repr_with_seed():
"""Test __repr__ with seed parameter."""
step = MockStep("test_step")
pipeline = RobotProcessor([step], seed=42)
repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep], seed=42)"
assert repr_str == expected
def test_repr_with_custom_name_and_seed():
"""Test __repr__ with both custom name and seed."""
step1 = MockStep("step1")
step2 = MockStepWithoutOptionalMethods()
pipeline = RobotProcessor([step1, step2], name="MyProcessor", seed=123)
repr_str = repr(pipeline)
expected = (
"RobotProcessor(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods], seed=123)"
)
assert repr_str == expected
def test_repr_without_seed():
"""Test __repr__ when seed is explicitly None (should not show seed)."""
step = MockStep("test_step")
pipeline = RobotProcessor([step], name="TestProcessor", seed=None)
repr_str = repr(pipeline)
expected = "RobotProcessor(name='TestProcessor', steps=1: [MockStep])"
assert repr_str == expected
def test_repr_various_step_types():
"""Test __repr__ with different types of steps to verify class name extraction."""
step1 = MockStep()
step2 = MockStepWithTensorState()
step3 = MockModuleStep()
step4 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4], name="MixedSteps")
repr_str = repr(pipeline)
expected = "RobotProcessor(name='MixedSteps', steps=4: [MockStep, MockStepWithTensorState, ..., MockNonModuleStepWithState])"
assert repr_str == expected
def test_repr_edge_case_long_names():
"""Test __repr__ handles steps with long class names properly."""
step1 = MockStepWithNonSerializableParam()
step2 = MockStepWithoutOptionalMethods()
step3 = MockStepWithTensorState()
step4 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4], name="LongNames", seed=999)
repr_str = repr(pipeline)
expected = "RobotProcessor(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState], seed=999)"
assert repr_str == expected