Extended HMSET

aioredis/commands/hash.py

@@ -1,4 +1,11 @@
-from aioredis.util import wait_ok, wait_convert, wait_make_dict, _NOTSET, PY_35
+from aioredis.util import (
+    wait_ok,
+    wait_convert,
+    wait_make_dict,
+    make_pairs,
+    _NOTSET,
+    PY_35,
+    )
 
 if PY_35:
     from aioredis.util import _ScanIterPairs
@@ -50,12 +57,10 @@ def hmget(self, key, field, *fields, encoding=_NOTSET):
         return self._conn.execute(b'HMGET', key, field, *fields,
                                   encoding=encoding)
 
-    # TODO: replace args with dict_or_pairs
-    def hmset(self, key, field, value, *pairs):
+    def hmset(self, key, *args, **kwargs):
         """Set multiple hash fields to multiple values."""
-        if len(pairs) % 2 != 0:
-            raise TypeError("length of pairs must be even number")
-        return wait_ok(self._conn.execute(b'HMSET', key, field, value, *pairs))
+        pairs = make_pairs(*args, **kwargs)
+        return wait_ok(self._conn.execute(b'HMSET', key, *pairs))
 
     def hset(self, key, field, value):
         """Set the string value of a hash field."""

aioredis/util.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import sys
+from itertools import chain
 
 from asyncio.base_events import BaseEventLoop
 
@@ -319,3 +320,22 @@ def create_future(loop):
 else:
     def create_future(loop):
         return asyncio.Future(loop=loop)
+
+
+def make_pairs(*args, **kwargs):
+    if not args and not kwargs:
+        raise TypeError("length of args or kwargs must be > 0")
+
+    elif len(args) == 1 and isinstance(args[0], dict):
+        pairs = chain.from_iterable(args[0].items())
+
+    elif len(args) % 2 != 0:
+        raise TypeError("length of args must be even number")
+
+    else:
+        pairs = args
+
+    if kwargs:
+        pairs = chain(pairs, chain.from_iterable(kwargs.items()))
+
+    return tuple(pairs)

tests/hash_commands_test.py

@@ -237,12 +237,44 @@ def test_hmset(redis):
     test_value = yield from redis.hmget(key, b'foo', b'bar')
     assert set(test_value) == {b'baz', b'paz'}
 
+    # dict
+    d1 = {b'foo': b'one dict'}
+    test_value = yield from redis.hmset(key, d1)
+    assert test_value is True
+    test_value = yield from redis.hget(key, b'foo')
+    assert test_value == b'one dict'
+
+    # kwdict
+    test_value = yield from redis.hmset(key, foo=b'kw1', bar=b'kw2')
+    assert test_value is True
+    test_value = yield from redis.hmget(key, b'foo', b'bar')
+    assert set(test_value) == {b'kw1', b'kw2'}
+
+    # pair & kwdict
+    test_value = yield from redis.hmset(key, b'foo', 'pairs', foo=b'kw1')
+    assert test_value is True
+    test_value = yield from redis.hget(key, b'foo')
+    assert test_value == b'kw1'
+
+    # dict & kwdict
+    d1 = {b'foo': b'dict'}
+    test_value = yield from redis.hmset(key, d1, foo=b'kw')
+    assert test_value is True
+    test_value = yield from redis.hget(key, b'foo')
+    assert test_value == b'kw'
+
+    with pytest.raises(TypeError):
+        yield from redis.hmset(key, {'a': 1}, {'b': 2}, 'c', 3, d=4)
+
     with pytest.raises(TypeError):
         yield from redis.hmset(key, b'foo', b'bar', b'baz')
 
     with pytest.raises(TypeError):
         yield from redis.hmset(None, *pairs)
 
+    with pytest.raises(TypeError):
+        yield from redis.hmset(key)
+
 
 @pytest.mark.run_loop
 def test_hset(redis):