python hier/top block holds refs of connected blocks

4 files changed, 39 insertions, 23 deletions

diff --git a/include/gras/element.i b/include/gras/element.i
index 364328c..0776362 100644
--- a/include/gras/element.i
+++ b/include/gras/element.i
@@ -42,8 +42,14 @@ namespace gras
     {
         std::string __str__(void)
         {
-           return self->to_string();
+           return ($self)->to_string();
         }
+
+        bool __eq__(const Element &rhs)
+        {
+            return ($self)->get() == rhs.get();
+        }
+
     }
 }
 
diff --git a/lib/element.cpp b/lib/element.cpp
index 6d116e9..83466c2 100644
--- a/lib/element.cpp
+++ b/lib/element.cpp
@@ -52,8 +52,12 @@ ElementImpl::~ElementImpl(void)
 
 Element &Element::shared_to_element(void)
 {
-    this->weak_self = this->shared_from_this();
-    return *this;
+    try
+    {
+        this->weak_self = this->shared_from_this();
+    }
+    catch(...){}
+    return this->derived_to_element();
 }
 
 long Element::unique_id(void) const
diff --git a/lib/hier_block.cpp b/lib/hier_block.cpp
index 36e1b9d..6c95e9b 100644
--- a/lib/hier_block.cpp
+++ b/lib/hier_block.cpp
@@ -45,13 +45,6 @@ void HierBlock::disconnect(const Element &elem)
     (*this)->topology->remove_topology(elem->topology.get());
 }
 
-static Apology::Wax get_elem_ref(const Element &elem)
-{
-    boost::shared_ptr<Element> shared_elem = elem.weak_self.lock();
-    if (shared_elem) return shared_elem;
-    return elem;
-}
-
 void HierBlock::connect(
     const Element &src,
     const size_t src_index,
@@ -60,8 +53,8 @@ void HierBlock::connect(
 ){
     //TODO, this is the perfect place to validate IO sigs
     const Apology::Flow flow(
-        Apology::Port(src->get_elem(), src_index, get_elem_ref(src)),
-        Apology::Port(sink->get_elem(), sink_index, get_elem_ref(sink))
+        Apology::Port(src->get_elem(), src_index, src.weak_self.lock()),
+        Apology::Port(sink->get_elem(), sink_index, sink.weak_self.lock())
     );
     (*this)->topology->add_flow(flow);
 }
diff --git a/python/gras/GRAS_HierBlock.i b/python/gras/GRAS_HierBlock.i
index 8fb972a..16142e5 100644
--- a/python/gras/GRAS_HierBlock.i
+++ b/python/gras/GRAS_HierBlock.i
@@ -109,42 +109,55 @@ struct HierBlockPython : HierBlock
 ////////////////////////////////////////////////////////////////////////
 %pythoncode %{
 
-def internal_connect__(fcn, obj, *args):
+def to_element(obj):
+    if isinstance(obj, Element): return obj
+    try: return obj.shared_to_element()
+    except: raise Exception('cant coerce obj %s to element'%(obj))
 
-    def to_element(obj):
-        if isinstance(obj, Element): return obj
-        try: return obj.shared_to_element()
-        except: raise Exception('cant coerce obj %s to element'%(obj))
+def internal_connect__(fcn, obj, refs, action, *args):
 
     if len(args) == 1:
-        fcn(obj, to_element(args[0]))
+        elem = (args[0])
+        fcn(obj, elem)
+        if elem != to_element(obj): action(refs, elem)
         return
 
     for src, sink in zip(args, args[1:]):
+
         try: src, src_index = src
         except: src_index = 0
+        src = to_element(src)
+
         try: sink, sink_index = sink
         except: sink_index = 0
-        fcn(obj, to_element(src), src_index, to_element(sink), sink_index)
+        sink = to_element(sink)
+
+        fcn(obj, src, src_index, sink, sink_index)
+
+        #incr/decr the python obj ref counts
+        if src != to_element(obj): action(refs, src)
+        if sink != to_element(obj): action(refs, sink)
 
 class TopBlock(TopBlockPython):
     def __init__(self, *args, **kwargs):
         TopBlockPython.__init__(self, *args, **kwargs)
+        self.__refs = list()
 
     def connect(self, *args):
-        return internal_connect__(TopBlockPython.connect, self, *args)
+        return internal_connect__(TopBlockPython.connect, self, self.__refs, list.append, *args)
 
     def disconnect(self, *args):
-        return internal_connect__(TopBlockPython.disconnect, self, *args)
+        return internal_connect__(TopBlockPython.disconnect, self, self.__refs, list.remove, *args)
 
 class HierBlock(HierBlockPython):
     def __init__(self, *args, **kwargs):
         HierBlockPython.__init__(self, *args, **kwargs)
+        self.__refs = list()
 
     def connect(self, *args):
-        return internal_connect__(HierBlockPython.connect, self, *args)
+        return internal_connect__(HierBlockPython.connect, self, self.__refs, list.append, *args)
 
     def disconnect(self, *args):
-        return internal_connect__(HierBlockPython.disconnect, self, *args)
+        return internal_connect__(HierBlockPython.disconnect, self, self.__refs, list.remove, *args)
 
 %}