diff --git a/doc/python/mechanisms.rst b/doc/python/mechanisms.rst index b55aa460a8604f180eb5e60bded2a30eeda5c184..d744ddfa933ffb214b464d37ba107d034429b196 100644 --- a/doc/python/mechanisms.rst +++ b/doc/python/mechanisms.rst @@ -246,10 +246,19 @@ Mechanism catalogues 2. A further hierarchy of *derived* mechanisms, that allow specialization of global parameters, ion bindings, and implementations. - .. py:method:: has(name) + .. py:method:: __contains__(name) Test if mechanism with *name* is in the catalogue. + Note: This enables the following idiom + + .. code-block:: Python + + import arbor + + if 'hh' in arbor.default_catalogue(): + print("Found HH mechanism.") + :param name: name of mechanism. :type name: str :return: bool @@ -280,11 +289,20 @@ Mechanism catalogues :return: mechanism metadata :rtype: :class:`mechanism_info` - .. py:method:: names() + .. py:method:: __iter___() Return a list names of all the mechanisms in the catalogue. - :return: list + Note: This enables the following idiom + + .. code-block:: Python + + import arbor + + for name in arbor.default_catalogue(): + print(name) + + :return: :class:`py_mech_cat_iterator` .. py:method:: derive(name, parent, globals={}, ions={}) diff --git a/python/mechanism.cpp b/python/mechanism.cpp index 70e2c5283b0ca88bb1791ce2f13ff526d1358b36..a1ac7a9d8f86120fa3d762dd8304f29b0193fc95 100644 --- a/python/mechanism.cpp +++ b/python/mechanism.cpp @@ -104,12 +104,29 @@ void register_mechanisms(pybind11::module& m) { return util::pprintf("(arbor.mechanism_info)"); }); pybind11::class_<arb::mechanism_catalogue> cat(m, "catalogue"); + + struct py_mech_cat_iterator { + py_mech_cat_iterator(const arb::mechanism_catalogue &cat, pybind11::object ref) : names(cat.mechanism_names()), ref(ref), idx{0} { } + std::vector<std::string> names; + pybind11::object ref; // keep a reference to cat lest it dies while we iterate + size_t idx = 0; + std::string next() { + if (idx == names.size()) throw pybind11::stop_iteration(); + return names[idx++]; + } + }; + + pybind11::class_<py_mech_cat_iterator>(cat, "MechCatIterator") + .def("__iter__", [](py_mech_cat_iterator &it) -> py_mech_cat_iterator& { return it; }) + .def("__next__", &py_mech_cat_iterator::next); + cat .def(pybind11::init<const arb::mechanism_catalogue&>()) - .def("has", &arb::mechanism_catalogue::has, + .def("__contains__", &arb::mechanism_catalogue::has, "name"_a, "Is 'name' in the catalogue?") - .def("keys", &arb::mechanism_catalogue::mechanism_names, - "Return a list of all mechanisms in this catalogues.") + .def("__iter__", + [](pybind11::object cat) { return py_mech_cat_iterator(cat.cast<const arb::mechanism_catalogue &>(), cat); }, + "Return an iterator over all mechanism names in this catalogues.") .def("is_derived", &arb::mechanism_catalogue::is_derived, "name"_a, "Is 'name' a derived mechanism or can it be implicitly derived?") .def("__getitem__",