Expose ContainerBase::childByIndex via operator[](int) (#413)

This commit is contained in:
Robert Haschke 2022-12-12 20:44:37 +01:00 committed by GitHub
parent 26c690c3b6
commit f872d58520
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 68 additions and 28 deletions

View File

@ -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

View File

@ -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<SerialContainer>("task pipeline"));

View File

@ -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) {

View File

@ -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()

View File

@ -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);
}