Add limit= to get methods, make their signatures more uniform

aetcd/__init__.py

@@ -22,6 +22,8 @@
 from .rtypes import KeyValue  # noqa: F401
 from .rtypes import Put  # noqa: F401
 from .rtypes import ResponseHeader  # noqa: F401
+from .rtypes import SortOrder  # noqa: F401
+from .rtypes import SortTarget  # noqa: F401
 from .rtypes import Watch  # noqa: F401
 from .watcher import Watcher  # noqa: F401
 from .watcher import WatcherCallback  # noqa: F401

aetcd/client.py

@@ -251,21 +251,36 @@ async def get(
     async def get_prefix(
         self,
         key_prefix: bytes,
-        sort_order: typing.Optional[str] = None,
-        sort_target: str = 'key',
+        limit: int = 0,
+        sort_order: rtypes.SortOrder = rtypes.SortOrder.NONE,
+        sort_target: rtypes.SortTarget = rtypes.SortTarget.KEY,
         keys_only: bool = False,
     ) -> rtypes.GetRange:
         """Get a range of keys with a prefix from the key-value store.
 
         :param bytes key_prefix:
             Key prefix to get.
 
+        :param aetcd.rtypes.SortOrder sort_order:
+            Order results in :class:`~aetcd.rtypes.SortOrder` direction.
+
+        :param aetcd.rtypes.SortTarget sort_target:
+            Order results by :class:`~aetcd.rtypes.SortTarget`.
+
+        :param int limit:
+            Limit on the number of keys returned for the request.
+            When limit is set to 0 (default), it is treated as no limit.
+
+        :param bool keys_only:
+            Returns only keys and not values.
+
         :return:
             An instance of :class:`~aetcd.rtypes.GetRange`.
         """
         range_request = self._build_get_range_request(
             key=key_prefix,
             range_end=utils.prefix_range_end(key_prefix),
+            limit=limit,
             sort_order=sort_order,
             sort_target=sort_target,
             keys_only=keys_only,
@@ -290,8 +305,10 @@ async def get_range(
         self,
         range_start: bytes,
         range_end: bytes,
-        sort_order: typing.Optional[str] = None,
-        sort_target: str = 'key',
+        limit: int = 0,
+        sort_order: rtypes.SortOrder = rtypes.SortOrder.NONE,
+        sort_target: rtypes.SortTarget = rtypes.SortTarget.KEY,
+        keys_only: bool = False,
     ) -> rtypes.GetRange:
         """Get a range of keys from the key-value store.
 
@@ -301,14 +318,29 @@ async def get_range(
         :param bytes range_end:
             Last key in range.
 
+        :param aetcd.rtypes.SortOrder sort_order:
+            Order results in :class:`~aetcd.rtypes.SortOrder` direction.
+
+        :param aetcd.rtypes.SortTarget sort_target:
+            Order results by :class:`~aetcd.rtypes.SortTarget`.
+
+        :param int limit:
+            Limit on the number of keys returned for the request.
+            When limit is set to 0 (default), it is treated as no limit.
+
+        :param bool keys_only:
+            Returns only keys and not values.
+
         :return:
             An instance of :class:`~aetcd.rtypes.GetRange`.
         """
         range_request = self._build_get_range_request(
             key=range_start,
             range_end=range_end,
+            limit=limit,
             sort_order=sort_order,
             sort_target=sort_target,
+            keys_only=keys_only,
         )
 
         range_response = await self.kvstub.Range(
@@ -328,18 +360,33 @@ async def get_range(
     @_ensure_connected
     async def get_all(
         self,
-        sort_order=None,
-        sort_target='key',
-        keys_only=False,
+        limit: int = 0,
+        sort_order: rtypes.SortOrder = rtypes.SortOrder.NONE,
+        sort_target: rtypes.SortTarget = rtypes.SortTarget.KEY,
+        keys_only: bool = False,
     ) -> rtypes.GetRange:
         """Get all keys from the key-value store.
 
+        :param aetcd.rtypes.SortOrder sort_order:
+            Order results in :class:`~aetcd.rtypes.SortOrder` direction.
+
+        :param aetcd.rtypes.SortTarget sort_target:
+            Order results by :class:`~aetcd.rtypes.SortTarget`.
+
+        :param int limit:
+            Limit on the number of keys returned for the request.
+            When limit is set to 0 (default), it is treated as no limit.
+
+        :param bool keys_only:
+            Returns only keys and not values.
+
         :return:
             An instance of :class:`~aetcd.rtypes.GetRange`.
         """
         range_request = self._build_get_range_request(
             key=b'\0',
             range_end=b'\0',
+            limit=limit,
             sort_order=sort_order,
             sort_target=sort_target,
             keys_only=keys_only,
@@ -1277,13 +1324,13 @@ async def snapshot(self, file_obj):
             file_obj.write(response.blob)
 
     @staticmethod
-    def _build_get_range_request(
+    def _build_get_range_request(  # noqa: C901
         key: bytes,
         range_end: typing.Optional[bytes] = None,
         limit: typing.Optional[int] = None,
         revision: typing.Optional[int] = None,
-        sort_order: typing.Optional[str] = None,
-        sort_target: str = 'key',
+        sort_order: rtypes.SortOrder = rtypes.SortOrder.NONE,
+        sort_target: rtypes.SortTarget = rtypes.SortTarget.KEY,
         serializable: bool = False,
         keys_only: bool = False,
         count_only: typing.Optional[int] = None,
@@ -1292,7 +1339,7 @@ def _build_get_range_request(
         min_create_revision: typing.Optional[int] = None,
         max_create_revision: typing.Optional[int] = None,
     ) -> rpc.RangeRequest:
-        # TODO: Add missing request parameters: limit, revision, count_only,
+        # TODO: Add missing request parameters: revision, count_only,
         #       mid_mod_revision, max_mod_revision, min_create_revision, max_create_revision
         range_request = rpc.RangeRequest()
 
@@ -1301,28 +1348,18 @@ def _build_get_range_request(
         if range_end is not None:
             range_request.range_end = range_end
 
-        if sort_order is None:
-            range_request.sort_order = rpc.RangeRequest.NONE
-        elif sort_order == 'ascend':
-            range_request.sort_order = rpc.RangeRequest.ASCEND
-        elif sort_order == 'descend':
-            range_request.sort_order = rpc.RangeRequest.DESCEND
-        else:
-            raise ValueError(f'unknown sort order: {sort_order!r}')
-
-        if sort_target is None or sort_target == 'key':
-            range_request.sort_target = rpc.RangeRequest.KEY
-        elif sort_target == 'version':
-            range_request.sort_target = rpc.RangeRequest.VERSION
-        elif sort_target == 'create':
-            range_request.sort_target = rpc.RangeRequest.CREATE
-        elif sort_target == 'mod':
-            range_request.sort_target = rpc.RangeRequest.MOD
-        elif sort_target == 'value':
-            range_request.sort_target = rpc.RangeRequest.VALUE
-        else:
-            raise ValueError('sort_target must be one of "key", '
-                             '"version", "create", "mod" or "value"')
+        if limit is not None:
+            range_request.limit = limit
+
+        range_request.sort_order = getattr(
+            rpc.RangeRequest.SortOrder,
+            rtypes.SortOrder(sort_order).name,
+        )
+
+        range_request.sort_target = getattr(
+            rpc.RangeRequest.SortTarget,
+            rtypes.SortTarget(sort_target).name,
+        )
 
         range_request.serializable = serializable
         range_request.keys_only = keys_only

aetcd/rtypes.py

@@ -15,6 +15,42 @@ def __repr__(self):
         ) + ']'
 
 
+class SortOrder(enum.Enum):
+    #: No sort order.
+    NONE = enum.auto()
+
+    #: Ascending sort order - lowest target value first.
+    ASCEND = enum.auto()
+
+    #: Descending sort order - highest target value first.
+    DESCEND = enum.auto()
+
+
+class SortTarget(enum.Enum):
+    #: Sort range operation results by key.
+    KEY = enum.auto()
+
+    #: Sort range operation results by key version.
+    VERSION = enum.auto()
+
+    #: Sort range operation results by revision of the last creation on the key.
+    CREATE = enum.auto()
+
+    #: Sort range operation results by revision of the last modification on the key.
+    MOD = enum.auto()
+
+    #: Sort range operation results by value.
+    VALUE = enum.auto()
+
+
+class EventKind(enum.Enum):
+    #: Designates a ``PUT`` event.
+    PUT = enum.auto()
+
+    #: Designates a ``DELETE`` event.
+    DELETE = enum.auto()
+
+
 class ResponseHeader(_Slotted):
     """Represents the metadata for the response."""
 
@@ -233,14 +269,6 @@ def __repr__(self):
         )
 
 
-class EventKind(str, enum.Enum):
-    #: Designates a ``PUT`` event.
-    PUT = 'PUT'
-
-    #: Designates a ``DELETE`` event.
-    DELETE = 'DELETE'
-
-
 class Event(_Slotted):
     """Reperesents a watch event."""
 
@@ -251,11 +279,11 @@ class Event(_Slotted):
     ]
 
     def __init__(self, kind, kv, prev_kv=None):
-        #: The kind of event. If the type is a ``PUT``, it indicates
-        #: new data has been stored to the key. If the type is a ``DELETE``,
-        #: it indicates the key was deleted.
-        self.kind: EventKind = rpc.Event.EventType.DESCRIPTOR.values_by_number[
-            kind].name
+        #: The kind of event.
+        #: ``PUT`` indicates that new data has been stored to the key.
+        #: ``DELETE`` indicates the key was deleted.
+        self.kind: EventKind = EventKind[
+            rpc.Event.EventType.DESCRIPTOR.values_by_number[kind].name]
 
         #: Holds the key-value for the event.
         #: A ``PUT`` event contains current key-value pair.

tests/integration/test_kv.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+import aetcd
 import aetcd.exceptions
 import aetcd.rpc
 
@@ -83,6 +84,19 @@ async def test_get_prefix_with_keys_only(etcdctl, etcd):
         assert not result.value
 
 
+@pytest.mark.asyncio
+async def test_get_prefix_with_limit(etcdctl, etcd):
+    for i in range(10):
+        etcdctl('put', f'/inrange{i}', 'in range')
+
+    results = list(await etcd.get_prefix(b'/inrange', limit=3))
+
+    assert len(results) == 3
+    for result in results:
+        assert result.key.startswith(b'/inrange')
+        assert result.value == b'in range'
+
+
 @pytest.mark.asyncio
 async def test_get_range(etcdctl, etcd):
     for char in string.ascii_lowercase:
@@ -98,6 +112,35 @@ async def test_get_range(etcdctl, etcd):
         assert result.value == b'in range'
 
 
+@pytest.mark.asyncio
+async def test_get_range_with_keys_only(etcdctl, etcd):
+    for i in range(20):
+        etcdctl('put', f'/inrange{i}', 'in range')
+
+    for i in range(5):
+        etcdctl('put', f'/notinrange{i}', 'not in range')
+
+    results = list(await etcd.get_range(b'/inrange', b'/inrangf', keys_only=True))
+
+    assert len(results) == 20
+    for result in results:
+        assert result.key.startswith(b'/inrange')
+        assert not result.value
+
+
+@pytest.mark.asyncio
+async def test_get_range_with_limit(etcdctl, etcd):
+    for i in range(10):
+        etcdctl('put', f'/inrange{i}', 'in range')
+
+    results = list(await etcd.get_range(b'/inrange', b'/inrangf', limit=3))
+
+    assert len(results) == 3
+    for result in results:
+        assert result.key.startswith(b'/inrange')
+        assert result.value == b'in range'
+
+
 @pytest.mark.asyncio
 async def test_get_range_with_sort_order(etcdctl, etcd):
     def remove_prefix(key, prefix):
@@ -110,15 +153,15 @@ def remove_prefix(key, prefix):
         etcdctl('put', f'/key{k}', v)
 
     keys = b''
-    for result in await etcd.get_prefix(b'/key', sort_order='ascend'):
+    for result in await etcd.get_prefix(b'/key', sort_order=aetcd.SortOrder.ASCEND):
         keys += remove_prefix(result.key, b'/key')
 
     assert keys == initial_keys.encode('utf-8')
 
     reverse_keys = b''
     for result in await etcd.get_prefix(
         b'/key',
-        sort_order='descend',
+        sort_order=aetcd.SortOrder.DESCEND,
     ):
         reverse_keys += remove_prefix(result.key, b'/key')
 
@@ -155,6 +198,32 @@ async def test_get_all_not_found(etcd):
     assert not result
 
 
+@pytest.mark.asyncio
+async def test_get_all_with_keys_only(etcdctl, etcd):
+    for i in range(10):
+        etcdctl('put', f'/inrange{i}', 'in range')
+
+    results = list(await etcd.get_all(keys_only=True))
+
+    assert len(results) == 10
+    for result in results:
+        assert result.key.startswith(b'/inrange')
+        assert not result.value
+
+
+@pytest.mark.asyncio
+async def test_get_all_with_limit(etcdctl, etcd):
+    for i in range(10):
+        etcdctl('put', f'/inrange{i}', 'in range')
+
+    results = list(await etcd.get_all(limit=3))
+
+    assert len(results) == 3
+    for result in results:
+        assert result.key.startswith(b'/inrange')
+        assert result.value == b'in range'
+
+
 @pytest.mark.asyncio
 async def test_put_key(etcdctl, etcd):
     await etcd.put(b'/key', b'value')