feat: add pyi for flatten method

functionalpy/_sequence.pyi

@@ -0,0 +1,65 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass
+from typing import Generic, TypeVar, overload
+
+_T0 = TypeVar("_T0")
+_T1 = TypeVar("_T1")
+_S = TypeVar("_S")
+@dataclass(frozen=True)
+class Group(Generic[_T0]):
+    key: str
+    value: Seq[_T0]
+    ...
+
+
+class Seq(Generic[_T0]):
+    def __init__(self, iterable: Iterable[_T0]) -> None:
+        ...
+
+    def count(self) -> int:
+        ...
+
+    def to_list(self) -> list[_T0]:
+        ...
+
+    def to_tuple(self) -> tuple[_T0, ...]:
+        ...
+
+    def to_iter(self) -> Iterator[_T0]:
+        ...
+
+    def to_set(self) -> set[_T0]:
+        ...
+
+    def map(self, func: Callable[[_T0], _T1]) -> Seq[_T1]:
+        ...
+
+    def filter(self, func: Callable[[_T0], bool]) -> Seq[_T0]:
+        ...
+
+    def reduce(self, func: Callable[[_T0, _T0], _T0]) -> _T0:
+        ...
+
+    def group_by(self, func: Callable[[_T0], str]) -> Seq[Group[_T0]]:
+        ...
+
+    @overload
+    def flatten(self: Seq[list[_S]]) -> Seq[_S]:
+        ...
+
+    @overload
+    def flatten(self: Seq[tuple[_S, ...]]) -> Seq[_S]:
+        ...
+
+    @overload
+    def flatten(self: Seq[_S]) -> Seq[_S]:
+        ...
+
+
+
+
+

functionalpy/test_sequence.py

@@ -1,4 +1,6 @@
-from functionalpy import Seq
+from cgi import test
+
+from functionalpy._sequence import Seq
 
 
 def test_chaining():
@@ -30,6 +32,26 @@ def test_reduce():
 
 
 def test_flatten():
-    sequence = Seq([[1, 2], [3, 4]])
-    result = sequence.flatten().to_list()
-    assert result == [1, 2, 3, 4]
+    test_input = ((1, 2), (3, 4))
+    sequence = Seq(test_input)
+    result = sequence.flatten()
+    assert result.to_list() == [1, 2, 3, 4]
+
+class TestFlattenTypes():
+    def test_flatten_tuple(self):
+        test_input = ((1, 2), (3, 4))
+        sequence = Seq(test_input)
+        result = sequence.flatten()
+        assert result.to_list() == [1, 2, 3, 4]
+
+    def test_flatten_list(self): 
+        test_input = [[1,2], [3,4]]
+        sequence = Seq(test_input)
+        result = sequence.flatten()
+        assert result.to_list() == [1, 2, 3, 4]
+
+    def test_flatten_str(self):
+        test_input = ['abcd']
+        sequence = Seq(test_input)
+        result = sequence.flatten()
+        assert result.to_list() == ['a', 'b', 'c', 'd']