From f872d585201ed3490b9e28591484351191d32102 Mon Sep 17 00:00:00 2001 From: Robert Haschke Date: Mon, 12 Dec 2022 20:44:37 +0100 Subject: [PATCH] Expose ContainerBase::childByIndex via operator[](int) (#413) --- .../moveit/task_constructor/container.h | 1 + core/include/moveit/task_constructor/task.h | 1 + core/python/bindings/src/core.cpp | 18 +++++ core/python/test/test_mtc.py | 70 +++++++++++-------- core/src/container.cpp | 6 ++ 5 files changed, 68 insertions(+), 28 deletions(-) diff --git a/core/include/moveit/task_constructor/container.h b/core/include/moveit/task_constructor/container.h index cea4a418..6f142e9a 100644 --- a/core/include/moveit/task_constructor/container.h +++ b/core/include/moveit/task_constructor/container.h @@ -53,6 +53,7 @@ public: size_t numChildren() const; Stage* findChild(const std::string& name) const; + Stage* operator[](int index) const; /** Callback function type used by traverse functions * Receives currently visited Stage and current depth in hierarchy diff --git a/core/include/moveit/task_constructor/task.h b/core/include/moveit/task_constructor/task.h index 54145abb..36ff5f94 100644 --- a/core/include/moveit/task_constructor/task.h +++ b/core/include/moveit/task_constructor/task.h @@ -74,6 +74,7 @@ class Task : protected WrapperBase public: PRIVATE_CLASS(Task) using WrapperBase::setCostTerm; + using WrapperBase::operator[]; Task(const std::string& ns = "", bool introspection = true, ContainerBase::pointer&& container = std::make_unique("task pipeline")); diff --git a/core/python/bindings/src/core.cpp b/core/python/bindings/src/core.cpp index c0500be0..d216e268 100644 --- a/core/python/bindings/src/core.cpp +++ b/core/python/bindings/src/core.cpp @@ -356,6 +356,15 @@ void export_core(pybind11::module& m) { return child; }, py::return_value_policy::reference_internal) + .def( + "__getitem__", + [](const ContainerBase& c, int idx) -> Stage* { + Stage* child = c[idx]; + if (!child) + throw py::index_error(); + return child; + }, + py::return_value_policy::reference_internal) .def( "__iter__", [](const ContainerBase& c) { @@ -441,6 +450,15 @@ void export_core(pybind11::module& m) { return child; }, py::return_value_policy::reference_internal) + .def( + "__getitem__", + [](const Task& t, int idx) -> Stage* { + Stage* child = t.stages()->operator[](idx); + if (!child) + throw py::index_error(); + return child; + }, + py::return_value_policy::reference_internal) .def( "__iter__", [](const Task& t) { diff --git a/core/python/test/test_mtc.py b/core/python/test/test_mtc.py index 07cd8be9..19cd843b 100644 --- a/core/python/test/test_mtc.py +++ b/core/python/test/test_mtc.py @@ -279,38 +279,52 @@ class TestStages(unittest.TestCase): stage.setCostTerm(costs) -class TestContainer(unittest.TestCase): +class BaseTestCases: + class ContainerTest(unittest.TestCase): + def __init__(self, ContainerType, *args, **kwargs): + super(BaseTestCases.ContainerTest, self).__init__(*args, **kwargs) + self.ContainerType = ContainerType + self.container = container = ContainerType() + container.add(stages.CurrentState("1")) + container.add(stages.CurrentState("2")) + container.add(stages.CurrentState("3")) + + def test_move(self): + container = self.ContainerType() + stage = stages.CurrentState() + container.add(stage) + with self.assertRaises(ValueError): + stage.name + + def test_access_by_name(self): + with self.assertRaises(IndexError): + self.container["unknown"] + + child = self.container["2"] + self.assertEqual(child.name, "2") + + def test_access_by_iterator(self): + self.assertEqual([child.name for child in self.container], ["1", "2", "3"]) + + def test_access_by_index(self): + self.assertEqual(self.container[0].name, "1") + self.assertEqual(self.container[1].name, "2") + self.assertEqual(self.container[-1].name, "3") + self.assertEqual(self.container[-2].name, "2") + with self.assertRaises(IndexError): + self.container[3] + with self.assertRaises(IndexError): + self.container[-4] + + +class TestSerial(BaseTestCases.ContainerTest): def __init__(self, *args, **kwargs): - super(TestContainer, self).__init__(*args, **kwargs) - - def test_move(self): - container = core.SerialContainer() - stage = stages.CurrentState() - container.add(stage) - with self.assertRaises(ValueError): - stage.name - - def check(self, container): - container = core.SerialContainer() - container.add(stages.CurrentState("1")) - container.add(stages.CurrentState("2")) - container.add(stages.CurrentState("3")) - with self.assertRaises(IndexError): - container["unknown"] - child = container["2"] - self.assertEqual(child.name, "2") - self.assertEqual([child.name for child in container], ["1", "2", "3"]) - - def test_serial(self): - self.check(core.SerialContainer()) - - def test_task(self): - self.check(core.Task()) + super(TestSerial, self).__init__(core.SerialContainer, *args, **kwargs) -class TestTask(unittest.TestCase): +class TestTask(BaseTestCases.ContainerTest): def __init__(self, *args, **kwargs): - super(TestTask, self).__init__(*args, **kwargs) + super(TestTask, self).__init__(core.Task, *args, **kwargs) def test(self): task = core.Task() diff --git a/core/src/container.cpp b/core/src/container.cpp index ace55373..b9043ebf 100644 --- a/core/src/container.cpp +++ b/core/src/container.cpp @@ -341,6 +341,12 @@ Stage* ContainerBase::findChild(const std::string& name) const { return nullptr; } +Stage* ContainerBase::operator[](int index) const { + auto impl = pimpl(); + auto it = impl->childByIndex(index, false); + return it != impl->children().end() ? it->get() : nullptr; +} + bool ContainerBase::traverseChildren(const ContainerBase::StageCallback& processor) const { return pimpl()->traverseStages(processor, 0, 1); }