ConnectorRelation improvements

- Use in export application - Add convenience methods
catmaid · Oct 31, 2018 · 47b2695 · 47b2695
1 parent f5cf5e4
commit 47b2695
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/catpy/client.py b/catpy/client.py
@@ -15,11 +15,19 @@
 
 
 class ConnectorRelation(IntEnum):
-    """Enum describing the type of a treenode-connector relationship, i.e. the treenode is ____ to the connector"""
+    """Enum describing the type of a treenode->connector relationship, i.e. the treenode is ____ to the connector"""
     PRESYNAPTIC_TO = 0
     POSTSYNAPTIC_TO = 1
     GAPJUNCTION_WITH = 2
     OTHER = -1
+    # todo: others?
+
+    @property
+    def is_synaptic(self):
+        return self == type(self).PRESYNAPTIC_TO or self == type(self).POSTSYNAPTIC_TO
+
+    def __str__(self):
+        return self.name.lower()
 
 
 class StackOrientation(IntEnum):
@@ -31,11 +39,7 @@ def __str__(self):
         return self.name.lower()
 
 
-orientation_strs = {
-    StackOrientation.XY: 'xy',
-    StackOrientation.XZ: 'xz',
-    StackOrientation.ZY: 'zy'
-}
+orientation_strs = {o: str(o) for o in StackOrientation}
 
 
 def make_url(base_url, *args):

diff --git a/catpy/export.py b/catpy/export.py
@@ -4,7 +4,7 @@
 
 from networkx.readwrite import json_graph
 
-from catpy.client import CatmaidClientApplication
+from catpy.client import CatmaidClientApplication, ConnectorRelation
 
 
 class ExportWidget(CatmaidClientApplication):
@@ -150,19 +150,19 @@ def get_treenode_and_connector_geometry(self, *skeleton_ids):
                 }
 
             for connector in data[1]:
-                if connector[2] not in [0, 1]:
+                relation = ConnectorRelation(connector[2])
+                if not relation.is_synaptic:
                     continue
 
                 conn_id = int(connector[1])
                 if conn_id not in skeleton['connectors']:
                     skeleton['connectors'][conn_id] = {
-                        'presynaptic_to': [],
-                        'postsynaptic_to': []
+                        str(r): [] for r in ConnectorRelation if r.is_synaptic
                     }
 
                 skeleton['connectors'][conn_id]['location'] = connector[3:6]
-                relation = 'postsynaptic_to' if connector[2] == 1 else 'presynaptic_to'
-                skeleton['connectors'][conn_id][relation].append(connector[0])
+                relation_name = str(relation)
+                skeleton['connectors'][conn_id][relation_name].append(connector[0])
 
             skeletons[int(skeleton_id)] = skeleton