From bc162de3ddcd30297724c305bcd4c34b9e38f674 Mon Sep 17 00:00:00 2001 From: Zhiyuan Liu Date: Fri, 17 Jul 2020 16:38:36 -0400 Subject: [PATCH] ENH: Add a convenient function to get/set the number of threads. --- Wrapping/Generators/Python/Tests/extras.py | 4 ++++ Wrapping/Generators/Python/itkExtras.py | 20 ++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/Wrapping/Generators/Python/Tests/extras.py b/Wrapping/Generators/Python/Tests/extras.py index 2e57493e528..9b72689f4cd 100644 --- a/Wrapping/Generators/Python/Tests/extras.py +++ b/Wrapping/Generators/Python/Tests/extras.py @@ -30,6 +30,10 @@ def custom_callback(name, progress): import sys import os +# test setting the number of threads +itk.set_nthreads(4) +assert itk.get_nthreads() == 4 + # test the force load function itk.force_load() diff --git a/Wrapping/Generators/Python/itkExtras.py b/Wrapping/Generators/Python/itkExtras.py index feb95d77df3..51a697fa556 100644 --- a/Wrapping/Generators/Python/itkExtras.py +++ b/Wrapping/Generators/Python/itkExtras.py @@ -27,6 +27,26 @@ # [K erases the end of the line clrLine = "\033[2000D\033[K" +def set_nthreads(number_of_threads): + """ + Support convenient set of the number of threads. + Use example (in python): + import itk + itk.set_nthreads(4) ## use 4 threads + """ + assert number_of_threads > 0, "Please set a possitive number of threads instead of %d" % number_of_threads + + import itk + threader = itk.MultiThreaderBase.New() + threader.SetGlobalDefaultNumberOfThreads(number_of_threads) + +def get_nthreads(): + """ + Get the number of threads + """ + import itk + threader = itk.MultiThreaderBase.New() + return threader.GetGlobalDefaultNumberOfThreads() def auto_not_in_place(v=True): """Force it to not run in place