import glob
import os
import sys
import defusedxml.ElementTree as ET

def write_swc(nmls_path, radius=0):
    node_radius = radius

    # store paths to nmls
    nmls = []
    if os.path.isdir(nmls_path):
        nmls = glob.glob(os.path.normpath(nmls_path) + '/*.nml')
    else:
        nmls.append(nmls_path)

    # create paths for resulting swcs
    swcs = []
    for nml in nmls:
        swcs.append(nml[:-4] + '.swc')
    for swc in swcs:
        if os.path.exists(swc):
            os.remove(swc)

    print('\nReading from .nml files...')
    nml_count = -1
    for nml in nmls:
        nml_count += 1
        tree = ET.parse(nml)
        thing_list = tree.findall('thing')
        for thing in thing_list:
            nodes = thing.find('nodes')
            edges = thing.find('edges')

            child_parent = []
            for edge in edges.findall('edge'):
                child = edge.get('target')
                parent = edge.get('source')
                child_parent.append((child, parent))

            child_list = [pair[0] for pair in child_parent]
            node_id, node_x, node_y, node_z, node_parent = ['']*5

            for node in nodes.findall('node'):
                node_id = node.get('id')
                node_x = node.get('x')
                node_y = node.get('y')
                node_z = float(node.get('z')) * 5.4545
                if node_radius == 0:
                    node_radius = node.get('radius')

                if node_id in child_list:
                    node_parent = child_parent[child_list.index(node_id)][1]
                else:
                    node_parent = -1

                swc = open(swcs[nml_count], 'a')
                swc.write(str(node_id) + ' 3 ' + str(node_x) + ' ' + str(node_y) + ' ' + str(node_z) + ' ' + str(node_radius) + ' ' + str(node_parent) + '\n')
                swc.close()

    # correct indexing: enforce consecutive natural numbering
    print('Writing final .swc files...')
    for swc in swcs:
        lines_to_write = []
        index_map = {'-1':-1}
        f = open(swc, 'r')
        lines = f.readlines()
        f.close()
        n = 1
        for line in lines:
            index_map[line.split(' ')[0]] = n
            n += 1

        for line in lines:
            line = str(index_map[line.split(' ')[0]]) + line[len(line.split(' ')[0]):]
            lines_to_write.append(line)

        n = 0
        for line in lines_to_write:
            line = line[:-len(line.split(' ')[6])] + str(index_map[(line.split(' ')[6])[:-1]])
            lines_to_write[n] = line
            n += 1
            if node_id in child_list:
                node_parent = child_parent[child_list.index(node_id)][1]
            else:
                node_parent = -1

        f = open(swc, 'w')
        for line in lines_to_write:
            f.write(line + '\n')
        f.close()
        print(swc)

if __name__ == "__main__":
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print('\nNML2SWC -- Written by Nathan Spencer 2016')
        print('Usage: python nml2swc.py ["path/to/nml/file.nml" || "path/to/nml/folder"] [radius]')
        print('Note: radius argument is optional; radius from nml will be used by default')
    elif len(sys.argv) == 2:
        write_swc(sys.argv[1])
    else:
        write_swc(sys.argv[1], sys.argv[2])

@@ -2011,23 +2011,275 @@
 write = 
-%5B%5D%0A    
+correct(swc, child_list, child_parent, node_id)%0A        f = open(swc, 'w')%0A        for line in lines_to_write:%0A            f.write(line + '%5Cn')%0A        f.close()%0A        print(swc)%0A%0Adef correct(swc, child_list, child_parent, node_id):%0A    lines_to_write = %5B%5D%0A
     inde
@@ -2296,28 +2296,24 @@
 -1':-1%7D%0A    
-  
-  
 f = open(swc
@@ -2319,20 +2319,16 @@
 c, 'r')%0A
-    
     line
@@ -2345,28 +2345,24 @@
 lines()%0A    
- 
-   
 f.close()%0A  
@@ -2355,28 +2355,24 @@
   f.close()%0A
-    
     n = 1%0A  
@@ -2365,28 +2365,24 @@
 )%0A    n = 1%0A
-    
     for line
@@ -2400,20 +2400,16 @@
         
-   
- 
 index_ma
@@ -2438,28 +2438,24 @@
 = n%0A        
-    
 n += 1%0A%0A    
@@ -2450,28 +2450,24 @@
 n += 1%0A%0A    
-  
-  
 for line in 
@@ -2469,28 +2469,24 @@
 e in lines:%0A
-    
         line
@@ -2556,36 +2556,32 @@
 )%5B0%5D):%5D%0A        
- 
-   
 lines_to_write.a
@@ -2593,20 +2593,16 @@
 (line)%0A%0A
-    
     n = 
@@ -2603,20 +2603,16 @@
   n = 0%0A
-    
     for 
@@ -2639,28 +2639,24 @@
 te:%0A        
-   
- 
 line = line%5B
@@ -2730,28 +2730,24 @@
 %5D%5D)%0A        
-    
 lines_to_wri
@@ -2771,23 +2771,15 @@
     
-   
- 
 n += 1%0A
-    
     
@@ -2812,36 +2812,32 @@
 st:%0A            
- 
-   
 node_parent = ch
@@ -2881,38 +2881,30 @@
 %5B1%5D%0A        
-    
 else:%0A
-    
             
@@ -2924,141 +2924,33 @@
  -1%0A
-%0A 
     
-   f = open(swc, 'w')%0A        for line in lines_to_write:%0A            f.write(line + '%5Cn')%0A        f.close()%0A        print(swc)
+return lines_to_write
 %0A%0Aif

split nml2swc into two functions [nml2swc]