diff --git a/catpy/client.py b/catpy/client.py index 9a88e9d..dfc9b9f 100644 --- 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 index 0c0b2c8..5492cec 100644 --- 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