From 299ad2bfe66a8ff5c2971119eaec571b2ac8bfc5 Mon Sep 17 00:00:00 2001
From: Chris Barnes <cbarnes@mrc-lmb.cam.ac.uk>
Date: Thu, 6 Jul 2023 14:02:34 +0100
Subject: [PATCH] IntEnum for treenode-connector relationships

---
 docs/source/whats_new.rst |  3 +++
 navis/core/__init__.py    |  5 +++--
 navis/core/base.py        | 28 ++++++++++++++++++++++++++--
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/docs/source/whats_new.rst b/docs/source/whats_new.rst
index 38dea986..b6e73563 100644
--- a/docs/source/whats_new.rst
+++ b/docs/source/whats_new.rst
@@ -18,6 +18,9 @@ repository.
    * - dev
      - ongoing
      - - dropped support for Python 3.7
+       - :class:`navis.NodeConnectorRelation` is an :class:`enum.IntEnum`
+         encoding relationships between (tree)nodes and connector nodes,
+         used in neurons' connector tables.
    * - 1.4.0
      - 21/12/22
      - - BREAKING: ``navis.flow_centrality`` was renamed to :func:`navis.synapse_flow_centrality`
diff --git a/navis/core/__init__.py b/navis/core/__init__.py
index 39b4342e..bf7bfb2a 100644
--- a/navis/core/__init__.py
+++ b/navis/core/__init__.py
@@ -12,7 +12,7 @@
 #    GNU General Public License for more details.
 
 from .volumes import Volume
-from .base import Neuron, BaseNeuron
+from .base import Neuron, BaseNeuron, NodeConnectorRelation
 from .skeleton import TreeNeuron
 from .mesh import MeshNeuron
 from .dotprop import Dotprops
@@ -25,4 +25,5 @@
 NeuronObject = Union[NeuronList, TreeNeuron, BaseNeuron, MeshNeuron]
 
 __all__ = ['Volume', 'Neuron', 'BaseNeuron', 'TreeNeuron', 'MeshNeuron',
-           'Dotprops', 'VoxelNeuron', 'NeuronList', 'make_dotprops']
+           'Dotprops', 'VoxelNeuron', 'NeuronList', 'make_dotprops',
+           'NodeConnectorRelation']
diff --git a/navis/core/base.py b/navis/core/base.py
index ca036448..83a0290c 100644
--- a/navis/core/base.py
+++ b/navis/core/base.py
@@ -12,6 +12,7 @@
 #    GNU General Public License for more details.
 
 import copy
+from enum import IntEnum
 import hashlib
 import numbers
 import pint
@@ -45,6 +46,27 @@
     pint.Quantity([])
 
 
+class NodeConnectorRelation(IntEnum):
+    """An integer describing a (tree)node-connector relationship.
+
+    i.e. "the (tree)node is <NodeConnectorType> the connector node"
+
+    Based on the `CATMAID link types`_.
+    A node PRESYNAPTIC_TO a connector is an output site.
+    A node POSTSYNAPTIC_TO a connector is an input site.
+
+    .. _`CATMAID link types`: https://github.com/catmaid/CATMAID/blob/2964e04e6e9772aff5d305e72c1b878030fe0e25/django/applications/catmaid/control/link.py#L16
+    """
+    PRESYNAPTIC_TO = 0
+    POSTSYNAPTIC_TO = 1
+    ABUTTING = 2
+    GAPJUNCTION_WITH = 3
+    TIGHTJUNCTION_WITH = 4
+    DESMOSOME_WITH = 5
+    ATTACHMENT_TO = 6
+    CLOSE_TO = 7
+
+
 def Neuron(x: Union[nx.DiGraph, str, pd.DataFrame, 'TreeNeuron', 'MeshNeuron'],
            **metadata):
     """Constructor for Neuron objects. Depending on the input, either a
@@ -421,7 +443,8 @@ def presynapses(self):
             raise ValueError('No connector table found.')
         # Make an educated guess what presynapses are
         types = self.connectors['type'].unique()
-        pre = [t for t in types if 'pre' in str(t) or t in [0, "0"]]
+        pre_int = NodeConnectorRelation.PRESYNAPTIC_TO.value
+        pre = [t for t in types if 'pre' in str(t) or t in [pre_int, str(pre_int)]]
 
         if len(pre) == 0:
             logger.debug(f'Unable to find presynapses in types: {types}')
@@ -442,7 +465,8 @@ def postsynapses(self):
             raise ValueError('No connector table found.')
         # Make an educated guess what presynapses are
         types = self.connectors['type'].unique()
-        post = [t for t in types if 'post' in str(t) or t in [1, "1"]]
+        post_int = NodeConnectorRelation.POSTSYNAPTIC_TO.value
+        post = [t for t in types if 'post' in str(t) or t in [post_int, str(post_int)]]
 
         if len(post) == 0:
             logger.debug(f'Unable to find postsynapses in types: {types}')