fix add_limit behavior in edge cases

dlt/extract/incremental/__init__.py

@@ -7,7 +7,6 @@
 from functools import wraps
 
 
-
 import dlt
 from dlt.common.exceptions import MissingDependencyException
 from dlt.common import pendulum, logger

dlt/extract/resource.py

@@ -313,8 +313,17 @@ def add_limit(self, max_items: int) -> "DltResource":  # noqa: A003
             "DltResource": returns self
         """
 
+        # make sure max_items is a number, to allow "None" as value for unlimited
+        if max_items == None:
+            max_items = -1
+
         def _gen_wrap(gen: TPipeStep) -> TPipeStep:
             """Wrap a generator to take the first `max_items` records"""
+
+            # zero items should produce empty generator
+            if max_items == 0:
+                return
+
             count = 0
             is_async_gen = False
             if inspect.isfunction(gen):

tests/extract/test_sources.py

@@ -2,6 +2,7 @@
 from typing import Iterator
 
 import pytest
+import asyncio
 
 import dlt
 from dlt.common.configuration.container import Container
@@ -789,6 +790,31 @@ def test_limit_infinite_counter() -> None:
     assert list(r) == list(range(10))
 
 
+@pytest.mark.parametrize("limit", (None, -1, 0, 10))
+def test_limit_edge_cases(limit: int) -> None:
+    r = dlt.resource(range(20), name="infinity").add_limit(limit)
+
+    @dlt.resource()
+    async def r_async():
+        for i in range(20):
+            await asyncio.sleep(0.01)
+            yield i
+
+    sync_list = list(r)
+    async_list = list(r_async().add_limit(limit))
+
+    # check the expected results
+    assert sync_list == async_list
+    if limit == 10:
+        assert sync_list == list(range(10))
+    elif limit in [None, -1]:
+        assert sync_list == list(range(20))
+    elif limit == 0:
+        assert sync_list == []
+    else:
+        assert False
+
+
 def test_limit_source() -> None:
     def mul_c(item):
         yield from "A" * (item + 2)