diff --git a/etcd3/client.py b/etcd3/client.py index 1971874b..ebc58b39 100644 --- a/etcd3/client.py +++ b/etcd3/client.py @@ -295,7 +295,7 @@ def get_prefix_response(self, key_prefix, **kwargs): range_request = self._build_get_range_request( key=key_prefix, - range_end=utils.increment_last_byte(utils.to_bytes(key_prefix)), + range_end=utils.prefix_range_end(utils.to_bytes(key_prefix)), **kwargs ) @@ -518,7 +518,7 @@ def delete_prefix(self, prefix): """Delete a range of keys with a prefix in etcd.""" delete_request = self._build_delete_request( prefix, - range_end=utils.increment_last_byte(utils.to_bytes(prefix)) + range_end=utils.prefix_range_end(utils.to_bytes(prefix)) ) return self.kvstub.DeleteRange( delete_request, @@ -586,7 +586,7 @@ def add_watch_prefix_callback(self, key_prefix, callback, **kwargs): :returns: watch_id. Later it could be used for cancelling watch. """ kwargs['range_end'] = \ - utils.increment_last_byte(utils.to_bytes(key_prefix)) + utils.prefix_range_end(utils.to_bytes(key_prefix)) return self.add_watch_callback(key_prefix, callback, **kwargs) @@ -665,7 +665,7 @@ def watch_prefix_response(self, key_prefix, **kwargs): :returns: tuple of ``responses_iterator`` and ``cancel``. """ kwargs['range_end'] = \ - utils.increment_last_byte(utils.to_bytes(key_prefix)) + utils.prefix_range_end(utils.to_bytes(key_prefix)) return self.watch_response(key_prefix, **kwargs) def watch_prefix(self, key_prefix, **kwargs): @@ -677,7 +677,7 @@ def watch_prefix(self, key_prefix, **kwargs): :returns: tuple of ``events_iterator`` and ``cancel``. """ kwargs['range_end'] = \ - utils.increment_last_byte(utils.to_bytes(key_prefix)) + utils.prefix_range_end(utils.to_bytes(key_prefix)) return self.watch(key_prefix, **kwargs) @_handle_errors @@ -730,7 +730,7 @@ def watch_prefix_once_response(self, key_prefix, timeout=None, **kwargs): will raise ``WatchTimedOut`` exception. """ kwargs['range_end'] = \ - utils.increment_last_byte(utils.to_bytes(key_prefix)) + utils.prefix_range_end(utils.to_bytes(key_prefix)) return self.watch_once_response(key_prefix, timeout=timeout, **kwargs) def watch_prefix_once(self, key_prefix, timeout=None, **kwargs): @@ -741,7 +741,7 @@ def watch_prefix_once(self, key_prefix, timeout=None, **kwargs): will raise ``WatchTimedOut`` exception. """ kwargs['range_end'] = \ - utils.increment_last_byte(utils.to_bytes(key_prefix)) + utils.prefix_range_end(utils.to_bytes(key_prefix)) return self.watch_once(key_prefix, timeout=timeout, **kwargs) @_handle_errors diff --git a/etcd3/utils.py b/etcd3/utils.py index 0a15d77f..5b9068eb 100644 --- a/etcd3/utils.py +++ b/etcd3/utils.py @@ -1,6 +1,10 @@ -def increment_last_byte(byte_string): - s = bytearray(byte_string) - s[-1] = s[-1] + 1 +def prefix_range_end(prefix): + """Create a bytestring that can be used as a range_end for a prefix.""" + s = bytearray(prefix) + for i in reversed(range(len(s))): + if s[i] < 0xff: + s[i] = s[i] + 1 + break return bytes(s) diff --git a/tests/test_etcd3.py b/tests/test_etcd3.py index 96aa9cf5..d573bfa2 100644 --- a/tests/test_etcd3.py +++ b/tests/test_etcd3.py @@ -555,7 +555,7 @@ def test_nested_transactions(self, etcd): def test_transaction_range_conditions(self, etcd): etcdctl('put', '/doot/key1', 'dootdoot') etcdctl('put', '/doot/key2', 'notdootdoot') - range_end = utils.increment_last_byte(utils.to_bytes('/doot/')) + range_end = utils.prefix_range_end(utils.to_bytes('/doot/')) compare = [etcd.transactions.value('/doot/', range_end) == 'dootdoot'] status, _ = etcd.transaction(compare=compare, success=[], failure=[]) assert not status @@ -980,8 +980,11 @@ def test_disarm_alarm(self, etcd): class TestUtils(object): - def test_increment_last_byte(self): - assert etcd3.utils.increment_last_byte(b'foo') == b'fop' + def test_prefix_range_end(self): + assert etcd3.utils.prefix_range_end(b'foo') == b'fop' + assert etcd3.utils.prefix_range_end(b'ab\xff') == b'ac\xff' + assert (etcd3.utils.prefix_range_end(b'a\xff\xff\xff\xff\xff') + == b'b\xff\xff\xff\xff\xff') def test_to_bytes(self): assert isinstance(etcd3.utils.to_bytes(b'doot'), bytes) is True diff --git a/tox.ini b/tox.ini index 3894492b..98f2b219 100644 --- a/tox.ini +++ b/tox.ini @@ -44,7 +44,7 @@ exclude = .venv,.git,.tox,dist,docs,*lib/python*,*egg,build,etcd3/etcdrpc/ application-import-names = etcd3 max-complexity = 10 # TODO add docstrings for public methods, modules, etc -ignore = D1 +ignore = D1, W503 [travis] python = 3.6: py36, flake8