Updated the return type

Author: 1227haran

data_structures/stacks/lexicographical_numbers.py

@@ -1,36 +1,37 @@
-def lexical_order(max_number: int) -> str:
+from typing import Iterator
+
+def lexical_order(max_number: int) -> Iterator[str]:
     """
     Generate numbers in lexical order from 1 to max_number.
 
-    >>> lexical_order(13)
+    >>> " ".join(lexical_order(13))
     '1 10 11 12 13 2 3 4 5 6 7 8 9'
-    >>> lexical_order(1)
-    '1'
-    >>> lexical_order(20)
+    >>> list(lexical_order(1))
+    ['1']
+    >>> " ".join(lexical_order(20))
     '1 10 11 12 13 14 15 16 17 18 19 2 20 3 4 5 6 7 8 9'
-    >>> lexical_order(25)
+    >>> " ".join(lexical_order(25))
     '1 10 11 12 13 14 15 16 17 18 19 2 20 21 22 23 24 25 3 4 5 6 7 8 9'
+    >>> list(lexical_order(12))
+    ['1', '10', '11', '12', '2', '3', '4', '5', '6', '7', '8', '9']
     """
 
-    ans = []
     stack = [1]
 
     while stack:
         num = stack.pop()
         if num > max_number:
             continue
 
-        ans.append(str(num))
+        yield(str(num))
         if (num % 10) != 9:
             stack.append(num + 1)
 
         stack.append(num * 10)
 
-    return " ".join(ans)
-
 
 if __name__ == "__main__":
     from doctest import testmod
 
     testmod()
-    print(f"Numbers from 1 to 25 in lexical order: {lexical_order(25)}")
+    print(f"Numbers from 1 to 25 in lexical order: {list(lexical_order(26))}")