Skip to content

Commit dce67ce

Browse files
anna-grimanna-grim
andauthored
feat: multithread swc read (#664)
Co-authored-by: anna-grim <anna.grim@alleninstitute.org>
1 parent a28d83b commit dce67ce

File tree

1 file changed

+16
-8
lines changed

1 file changed

+16
-8
lines changed

src/neuron_proofreader/utils/swc_util.py

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -236,13 +236,21 @@ def read_from_zip(self, zip_path):
236236
List of dictionaries whose keys and values are the attribute
237237
names and values from an SWC file.
238238
"""
239-
swc_dicts = deque()
240-
with ZipFile(zip_path, "r") as zip_file:
241-
swc_files = [f for f in zip_file.namelist() if f.endswith(".swc")]
242-
for f in swc_files:
243-
swc_dict = self.read_from_zipped_file(zip_file, f)
244-
if swc_dict:
245-
swc_dicts.append(swc_dict)
239+
with ThreadPoolExecutor() as executor:
240+
with ZipFile(zip_path, "r") as zf:
241+
# Submit threads
242+
threads = list()
243+
for f in [f for f in zf.namelist() if f.endswith(".swc")]:
244+
threads.append(
245+
executor.submit(self.read_from_zipped_file, zf, f)
246+
)
247+
248+
# Store results
249+
swc_dicts = deque()
250+
for thread in as_completed(threads):
251+
swc_dict = thread.result()
252+
if swc_dict:
253+
swc_dicts.append(swc_dict)
246254
return swc_dicts
247255

248256
def read_from_zipped_file(self, zip_file, path):
@@ -607,7 +615,7 @@ def get_segment_id(swc_name):
607615
Parameters
608616
----------
609617
swc_name : str
610-
SWC filename, expected to be in the format "<segment_id>.swc".
618+
SWC filename, expected to be in the format "{segment_id}.swc".
611619
612620
Returns
613621
-------

0 commit comments

Comments
 (0)