Skip to content

Commit a8f3fe9

Browse files
added type hints and fixed potential bug
Signed-off-by: Lukas Folle <lukas.folle@snke.com>
1 parent 416584d commit a8f3fe9

File tree

1 file changed

+14
-7
lines changed

1 file changed

+14
-7
lines changed

monai/transforms/utility/array.py

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1973,10 +1973,17 @@ class FlattenSequence(Transform, ReduceTrait):
19731973
def __init__(self):
19741974
super().__init__()
19751975

1976-
def __call__(self, data):
1977-
if isinstance(data, (list, tuple)):
1978-
if len(data) == 0:
1979-
return data
1980-
if isinstance(data[0], (list, tuple)):
1981-
return [item for sublist in data for item in sublist]
1982-
return data
1976+
def __call__(self, data: list | tuple | Any) -> list | tuple | Any:
1977+
"""
1978+
Flatten a nested sequence by one level.
1979+
Args:
1980+
data: Input data, can be a nested sequence.
1981+
Returns:
1982+
Flattened list if input is a nested sequence, otherwise returns data unchanged.
1983+
"""
1984+
if isinstance(data, (list, tuple)):
1985+
if len(data) == 0:
1986+
return data
1987+
if all(isinstance(item, (list, tuple)) for item in data):
1988+
return [item for sublist in data for item in sublist]
1989+
return data

0 commit comments

Comments
 (0)