util: add duck_check_type_is_iterator

kioss/_exec.py

@@ -13,6 +13,8 @@
 T = TypeVar("T")
 R = TypeVar("R")
 
+from kioss import _util
+
 
 class IteratorWrapper(Iterator[T]):
     def __init__(self, iterator: Iterator[T]):
@@ -29,22 +31,18 @@ def __init__(self, iterator: Iterator[Iterator[R]]) -> None:
 
     @staticmethod
     def _sanitize_input(expected_iterator_elem):
-        if not isinstance(expected_iterator_elem, Iterator):
-            raise TypeError(
-                f"Flattened elements must be iterators, but got {type(expected_iterator_elem)}"
-            )
+        _util.duck_check_type_is_iterator(expected_iterator_elem)
         return expected_iterator_elem
 
     def __next__(self) -> R:
         try:
-            return next(
-                FlatteningIteratorWrapper._sanitize_input(self.current_iterator_elem)
-            )
+            _util.duck_check_type_is_iterator(self.current_iterator_elem)
+            return next(self.current_iterator_elem)
         except StopIteration:
             while True:
-                self.current_iterator_elem = FlatteningIteratorWrapper._sanitize_input(
-                    super().__next__()
-                )
+                elem = super().__next__()
+                _util.duck_check_type_is_iterator(elem)
+                self.current_iterator_elem = elem
                 try:
                     return next(self.current_iterator_elem)
                 except StopIteration:

kioss/_plan.py

@@ -241,12 +241,11 @@ def __iter__(self) -> Iterator[T]:
         iterator = self.source()
         try:
             # duck-type checks that the object returned by the source is an iterator
-            iterator.__iter__
-            iterator.__next__
-        except AttributeError:
+            _util.duck_check_type_is_iterator(iterator)
+        except TypeError as e:
             raise TypeError(
                 f"source must be a callable returning an iterator (implements __iter__ and __next__ methods), but the object resulting from a call to source() was not an iterator: got '{iterator}' of type {type(iterator)}."
-            )
+            ) from e
         return iterator
 
 

kioss/_util.py

@@ -31,3 +31,27 @@ def iterate(it: Union[Iterator[T], Iterable[T]]) -> None:
 
 def identity(obj: T) -> T:
     return obj
+
+def duck_check_type_is_iterator(expected_iterator: Any) -> None:
+    """
+    Raises:
+        TypeError: If the expected_iterator does not implement __iter__ and __next__ methods.
+    """
+
+    try:
+        expected_iterator.__iter__
+        implements__iter__ = True
+    except AttributeError:
+        implements__iter__ = False
+    try:
+        expected_iterator.__next__
+        implements__next__ = True
+    except AttributeError:
+        implements__next__ = False
+
+    if not implements__iter__ and not implements__next__:
+        raise TypeError(f"Provided object is not an iterator because it does not implement __next__ and __iter__ methods")
+    if not implements__iter__:
+        raise TypeError(f"Provided object is not an iterator because it implements the __next__ but not the __iter__ one.")
+    if not implements__next__:
+        raise TypeError(f"Provided object is not an iterator because it implements the __iter__ but not the __next__ one.")