From d2623ecf991a8e3bddc79958ed755c78d82abfd8 Mon Sep 17 00:00:00 2001
From: thorstenhater <24411438+thorstenhater@users.noreply.github.com>
Date: Mon, 11 Jan 2021 13:05:12 +0100
Subject: [PATCH] More pythonic membership on mechanism_catalogue. (#1306)
Introduce two minor changes to the Python API
to handle mechanism_catalogues idiomatically.
Instead of
import arbor as A
cat = A.default_catalogue()
if cat.has('hh'):
print("Found HH.")
for mech in cat.keys():
print("*", mech)
we can now write
import arbor as A
cat = A.default_catalogue()
if 'hh' in cat:
print("Found HH.")
for mech in cat:
print("*", mech)
which is closer to the expectations of Python users.
---
doc/python/mechanisms.rst | 24 +++++++++++++++++++++---
python/mechanism.cpp | 23 ++++++++++++++++++++---
2 files changed, 41 insertions(+), 6 deletions(-)
diff --git a/doc/python/mechanisms.rst b/doc/python/mechanisms.rst
index b55aa460..d744ddfa 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 70e2c528..a1ac7a9d 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__",
--
GitLab