diff --git a/Wrapping/Generators/Python/Tests/extras.py b/Wrapping/Generators/Python/Tests/extras.py index 5495167acfc..1247b3cdc89 100644 --- a/Wrapping/Generators/Python/Tests/extras.py +++ b/Wrapping/Generators/Python/Tests/extras.py @@ -348,6 +348,9 @@ def custom_callback(name, progress): transform_dict = itk.dict_from_transform(transforms) transform_back = itk.transform_from_dict(transform_dict) +# Write single transform +itk.transformwrite(transforms[0], sys.argv[7], compression=True) + # pipeline, auto_pipeline and templated class are tested in other files # BridgeNumPy diff --git a/Wrapping/Generators/Python/itk/support/extras.py b/Wrapping/Generators/Python/itk/support/extras.py index f09e5232649..9c5741aa49d 100644 --- a/Wrapping/Generators/Python/itk/support/extras.py +++ b/Wrapping/Generators/Python/itk/support/extras.py @@ -1453,7 +1453,7 @@ def transformread(filename: fileiotype) -> List["itkt.TransformBase"]: def transformwrite( - transforms: List["itkt.TransformBase"], + transforms: Union["itkt.TransformBase", List["itkt.TransformBase"]], filename: fileiotype, compression: bool = False, ) -> None: @@ -1462,8 +1462,8 @@ def transformwrite( Parameters ---------- - transforms: list of itk.TransformBaseTemplate[itk.D] - Python list of the transforms to write. + transforms: single transform or list of transforms + Single transform Python list of the itk.TransformBaseTemplate[itk.D]'s to write. filename: Path to the transform file (typically a .h5 file). @@ -1476,6 +1476,8 @@ def transformwrite( writer = itk.TransformFileWriterTemplate[itk.D].New() writer.SetFileName(f"{filename}") writer.SetUseCompression(compression) + if not isinstance(transforms, list): + transforms = [transforms] for transform in transforms: writer.AddTransform(transform) writer.Update()