From 0a94fac6fcbfdc91d449f46d4e6609274b31cb74 Mon Sep 17 00:00:00 2001
From: Chris Barnes <cbarnes@mrc-lmb.cam.ac.uk>
Date: Thu, 9 Mar 2017 17:44:12 -0500
Subject: [PATCH] Add method replicating skeleton export

Method to return a dictionary of the same format as the 'Treenode and
connector geometry' option from the skeleton export widget.
---
 catmaid_interface.py | 49 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/catmaid_interface.py b/catmaid_interface.py
index 0974d08..0f28e48 100644
--- a/catmaid_interface.py
+++ b/catmaid_interface.py
@@ -79,6 +79,55 @@ def get_skeleton_json( skeleton_id, proj_opts, withtags = True):
         d.append([])
     return d
 
+def get_treenode_and_connector_geometry(skeleton_id, proj_opts):
+    """
+    See CATMAID code for original js implementation:
+    http://github.com/catmaid/CATMAID/blob/master/django/applications/catmaid/static/js/widgets/export-widget.js#L449
+
+    Parameters
+    ----------
+    skeleton_id
+    proj_opts
+
+    Returns
+    -------
+    dict
+        dict of the same form as the JSON file used by the 'Treenode and connector geometry' option in the CATMAID
+        Export widget.
+    """
+    url = proj_opts['baseurl'] + '/{}/{}/1/0/compact-skeleton'.format( proj_opts['project_id'], skeleton_id)
+    data = requests.get(
+        url, auth=catmaid_auth_token(proj_opts['token'], proj_opts['authname'], proj_opts['authpass'])
+    ).json()
+
+    skeleton = {
+        'treenodes': dict(),
+        'connectors': dict()
+    }
+
+    for treenode in data[0]:
+        skeleton['treenodes'][treenode[0]] = {
+            'location': treenode[3:6],
+            'parent_id': treenode[1]
+        }
+
+    for connector in data[1]:
+        if connector[2] not in [0, 1]:
+            continue
+
+        conn_id = connector[1]
+        if conn_id not in skeleton['connectors']:
+            skeleton['connectors'][conn_id] = {
+                'presynaptic_to': [],
+                'postsynaptic_to': []
+            }
+
+        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])
+
+    return skeleton
+
 def get_connector_info( connector_id, proj_opts):
     url = proj_opts['baseurl'] + '/{}/connectors/{}/'.format( proj_opts['project_id'], connector_id )
     d = requests.get( url, auth = catmaid_auth_token( proj_opts['token'], proj_opts['authname'], proj_opts['authpass'] ) ).json()