Meaningful Python types

src/nrnpython/nrnpy_hoc.cpp

@@ -11,6 +11,8 @@
 #include "nrnpy_utils.h"
 #include "../nrniv/shapeplt.h"
 #include <vector>
+#include <unordered_map>
+
 #include "nrnwrap_dlfcn.h"
 
 #if defined(NRNPYTHON_DYNAMICLOAD) && NRNPYTHON_DYNAMICLOAD > 0
@@ -101,6 +103,10 @@ static cTemplate* hoc_vec_template_;
 static cTemplate* hoc_list_template_;
 static cTemplate* hoc_sectionlist_template_;
 
+static std::vector<const char*> class_name_list;
+static std::unordered_map<Symbol*, PyTypeObject*> sym_to_type_map;
+static std::unordered_map<PyTypeObject*, Symbol*> type_to_sym_map;
+
 // typestr returned by Vector.__array_interface__
 // byteorder (first element) is modified at import time
 // to reflect the system byteorder
@@ -157,8 +163,13 @@ static PyObject* get_mech_object_ = NULL;
 static PyObject* nrnpy_rvp_pyobj_callback = NULL;
 
 PyTypeObject* hocobject_type;
+
 static PyObject* hocobj_call(PyHocObject* self, PyObject* args, PyObject* kwrds);
 
+void nrnpy_register_class(const char* name) {
+    class_name_list.push_back(name);
+}
+
 static PyObject* nrnexec(PyObject* self, PyObject* args) {
     const char* cmd;
     if (!PyArg_ParseTuple(args, "s", &cmd)) {
@@ -206,8 +217,12 @@ static void hocobj_dealloc(PyHocObject* self) {
 
 static PyObject* hocobj_new(PyTypeObject* subtype, PyObject* args, PyObject* kwds) {
     PyObject* subself;
+    PyObject* base;
+    PyHocObject* hbase;
+    bool has_base = false;
+    bool ok = true;
     subself = subtype->tp_alloc(subtype, 0);
-    // printf("hocobj_new %s %p\n", subtype->tp_name, subself);
+    //printf("hocobj_new %s %p %p\n", subtype->tp_name, subtype, subself);
     if (subself == NULL) {
         return NULL;
     }
@@ -219,35 +234,70 @@ static PyObject* hocobj_new(PyTypeObject* subtype, PyObject* args, PyObject* kwd
     self->nindex_ = 0;
     self->type_ = PyHoc::HocTopLevelInterpreter;
     self->iteritem_ = 0;
+
+    // if subtype is a subclass of some NEURON class, then one of its
+    // tp_mro's is in sym_to_type_map
+
+    for (Py_ssize_t i = 0; i < PyTuple_Size(subtype->tp_mro); i++) {
+        PyObject* item = PyTuple_GetItem(subtype->tp_mro, i);
+        Py_INCREF(item);
+        auto location = type_to_sym_map.find((PyTypeObject*) item);
+        if (location != type_to_sym_map.end()) {
+            PyObject* result = hocobj_new(hocobject_type, 0, 0);
+            hbase = (PyHocObject*) result;
+            hbase->type_ = PyHoc::HocFunction;
+            hbase->sym_ = location->second; //type_to_sym_map[(PyTypeObject*) item];
+            has_base = true;
+        }
+        Py_DECREF(item);
+        if (has_base) {
+            break;
+        }
+    }
+/*
+    if (PyObject_IsSubclass((PyObject*) subtype, (PyObject*) vectorobject_type)) {
+        //printf("subtype is vectorobject_type\n");
+        PyObject* result = hocobj_new(hocobject_type, 0, 0);
+        hbase = (PyHocObject*) result;
+        hbase->type_ = PyHoc::HocFunction;
+        hbase->sym_ = hoc_vec_template_->sym;
+        has_base = true;
+    } else {
+        //printf("not vectorobject_type\n");
+    }*/
     if (kwds && PyDict_Check(kwds)) {
-        PyObject* base = PyDict_GetItemString(kwds, "hocbase");
+        base = PyDict_GetItemString(kwds, "hocbase");
         if (base) {
-            int ok = 0;
+            ok = false;
             if (PyObject_TypeCheck(base, hocobject_type)) {
-                PyHocObject* hbase = (PyHocObject*) base;
-                if (hbase->type_ == PyHoc::HocFunction && hbase->sym_->type == TEMPLATE) {
-                    // printf("hocobj_new base %s\n", hbase->sym_->name);
-                    // remove the hocbase keyword since hocobj_call only allows
-                    // the "sec" keyword argument
-                    PyDict_DelItemString(kwds, "hocbase");
-                    PyObject* r = hocobj_call(hbase, args, kwds);
-                    if (!r) {
-                        Py_DECREF(subself);
-                        return NULL;
-                    }
-                    PyHocObject* rh = (PyHocObject*) r;
-                    self->type_ = rh->type_;
-                    self->ho_ = rh->ho_;
-                    hoc_obj_ref(self->ho_);
-                    Py_DECREF(r);
-                    ok = 1;
-                }
+                hbase = (PyHocObject*) base;
+                has_base = true;
             }
-            if (!ok) {
+            // TODO: is it a problem doing this before we're done with the base
+            PyDict_DelItemString(kwds, "hocbase");
+        }
+    }
+    if (has_base) {
+        if (hbase->type_ == PyHoc::HocFunction && hbase->sym_->type == TEMPLATE) {
+            // printf("hocobj_new base %s\n", hbase->sym_->name);
+            // remove the hocbase keyword since hocobj_call only allows
+            // the "sec" keyword argument
+            PyObject* r = hocobj_call(hbase, args, kwds);
+            if (!r) {
                 Py_DECREF(subself);
-                PyErr_SetString(PyExc_TypeError, "HOC base class not valid");
                 return NULL;
             }
+            PyHocObject* rh = (PyHocObject*) r;
+            self->type_ = rh->type_;
+            self->ho_ = rh->ho_;
+            hoc_obj_ref(self->ho_);
+            Py_DECREF(r);
+            ok = 1;
+        }
+        if (!ok) {
+            Py_DECREF(subself);
+            PyErr_SetString(PyExc_TypeError, "HOC base class not valid");
+            return NULL;
         }
     }
     return subself;
@@ -503,15 +553,22 @@ PyObject* nrnpy_ho2po(Object* o) {
     // The return value is None, or the encapsulated PyObject or
     // an encapsulating PyHocObject
     PyObject* po;
+    //printf("inside ho2po\n");
     if (!o) {
         po = Py_BuildValue("");
+        //printf("empty po\n");
     } else if (o->ctemplate->sym == nrnpy_pyobj_sym_) {
         po = nrnpy_hoc2pyobject(o);
         Py_INCREF(po);
+        //printf("from nrnpy_hoc2pyobject\n");
     } else {
         po = hocobj_new(hocobject_type, 0, 0);
         ((PyHocObject*) po)->ho_ = o;
         ((PyHocObject*) po)->type_ = PyHoc::HocObject;
+        auto location = sym_to_type_map.find(o->ctemplate->sym);
+        if (location != sym_to_type_map.end()) {
+            po->ob_type = sym_to_type_map[o->ctemplate->sym];
+        }
         hoc_obj_ref(o);
     }
     return po;
@@ -672,10 +729,19 @@ static void* fcall(void* vself, void* vargs) {
         double d = hoc_call_func(self->sym_, 1);
         hoc_pushx(d);
     } else if (self->sym_->type == TEMPLATE) {
+        //printf("in fcall TEMPLATE case\n");
         Object* ho = hoc_newobj1(self->sym_, narg);
         PyHocObject* result = (PyHocObject*) hocobj_new(hocobject_type, 0, 0);
         result->ho_ = ho;
         result->type_ = PyHoc::HocObject;
+        // Note: I think the only reason we're not using ho2po here is because we don't have to
+        //       hocref ho since it was created by hoc_newobj1... but it would be better if we did
+        //       so we could avoid repetitive code
+        auto location = sym_to_type_map.find(ho->ctemplate->sym);
+        if (location != sym_to_type_map.end()) {
+            ((PyObject*) result)->ob_type = location->second;
+        }
+
         hocobj_pushargs_free_strings(strings_to_free);
         return result;
     } else {
@@ -958,6 +1024,13 @@ static PyObject* hocobj_getattr(PyObject* subself, PyObject* pyname) {
     }
 
     Symbol* sym = getsym(n, self->ho_, 0);
+    // Return well known types right away
+    auto location = sym_to_type_map.find(sym);
+    if (location != sym_to_type_map.end()) {
+        Py_INCREF(location->second);
+        return (PyObject*) location->second;
+    }
+
     if (!sym) {
         if (self->type_ == PyHoc::HocObject && self->ho_->ctemplate->sym == nrnpy_pyobj_sym_) {
             PyObject* p = nrnpy_hoc2pyobject(self->ho_);
@@ -3074,8 +3147,22 @@ static char* nrncore_arg(double tstop) {
     return NULL;
 }
 
+
+static PyType_Spec obj_spec_from_name(const char* name) {
+    return {
+        name,
+        sizeof(PyHocObject),
+        0,
+        Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+        nrnpy_HocObjectType_slots,
+    };
+}
+
 PyObject* nrnpy_hoc() {
     PyObject* m;
+    PyObject* bases;
+    PyTypeObject* pto;
+    PyType_Spec spec;
     nrnpy_vec_from_python_p_ = nrnpy_vec_from_python;
     nrnpy_vec_to_python_p_ = nrnpy_vec_to_python;
     nrnpy_vec_as_numpy_helper_ = vec_as_numpy_helper;
@@ -3102,13 +3189,34 @@ PyObject* nrnpy_hoc() {
     m = PyModule_Create(&hocmodule);
     assert(m);
     Symbol* s = NULL;
-    hocobject_type = (PyTypeObject*) PyType_FromSpec(&nrnpy_HocObjectType_spec);
+    spec = obj_spec_from_name("hoc.HocObject");
+    hocobject_type = (PyTypeObject*) PyType_FromSpec(&spec);
     if (PyType_Ready(hocobject_type) < 0)
         goto fail;
     Py_INCREF(hocobject_type);
-    // printf("AddObject HocObject\n");
     PyModule_AddObject(m, "HocObject", (PyObject*) hocobject_type);
 
+    bases = PyTuple_Pack(1, hocobject_type);
+    Py_INCREF(bases);
+    for (auto name: class_name_list) {
+        // TODO: obj_spec_from_name needs a hoc. prepended
+        auto long_name = std::string("hoc.") + name;
+        spec = obj_spec_from_name(long_name.c_str());
+        pto = (PyTypeObject*) PyType_FromSpecWithBases(&spec, bases);
+        sym_to_type_map[hoc_lookup(name)] = pto;
+        type_to_sym_map[pto] = hoc_lookup(name);
+        if (PyType_Ready(pto) < 0)
+            goto fail;
+        Py_INCREF(pto);
+        PyModule_AddObject(m, name, (PyObject*) pto);
+    }
+    Py_DECREF(bases);
+
+    //printf("done\n");
+    // printf("AddObject HocObject\n");
+    //printf("registering Vector\n");
+    //printf("done\n");
+
     topmethdict = PyDict_New();
     for (PyMethodDef* meth = toplevel_methods; meth->ml_name != NULL; meth++) {
         PyObject* descr;

src/nrnpython/nrnpy_hoc.h

@@ -27,13 +27,6 @@ static PyType_Slot nrnpy_HocObjectType_slots[] = {
     {Py_nb_true_divide, (PyObject*) py_hocobj_div},
     {0, 0},
 };
-static PyType_Spec nrnpy_HocObjectType_spec = {
-    "hoc.HocObject",
-    sizeof(PyHocObject),
-    0,
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
-    nrnpy_HocObjectType_slots,
-};
 
 
 static struct PyModuleDef hocmodule = {PyModuleDef_HEAD_INIT,

src/oc/hoc_oop.cpp

@@ -30,6 +30,7 @@ Symbol* nrnpy_pyobj_sym_;
 void (*nrnpy_py2n_component)(Object* o, Symbol* s, int nindex, int isfunc);
 void (*nrnpy_hpoasgn)(Object* o, int type);
 void* (*nrnpy_opaque_obj2pyobj_p_)(Object*);
+extern void nrnpy_register_class(const char* name);
 #endif
 
 #if CABLE
@@ -1611,6 +1612,9 @@ void class2oc(const char* name,
     if (hoc_lookup(name)) {
         hoc_execerror(name, "already being used as a name");
     }
+    #if USE_PYTHON
+    nrnpy_register_class(name);
+    #endif
     tsym = hoc_install(name, UNDEF, 0.0, &hoc_symlist);
     tsym->subtype = CPLUSOBJECT;
     hoc_begintemplate(tsym);

test/cover/checkresult.py

@@ -1,6 +1,6 @@
 import json
 import math
-from neuron import h
+from neuron import h, hoc
 
 
 class Chk:
@@ -39,7 +39,7 @@ def __call__(self, key, value, tol=0.0):
         """
 
         if key in self.d:
-            if type(value) == type(h.Vector):  # actually hoc.HocObject
+            if isinstance(value, hoc.Vector):
                 # Convert to list to keep the `equal` method below simple
                 value = list(value)
             # Hand-rolled comparison that uses `tol` for arithmetic values