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