1from aioredis.util import wait_convert, _NOTSET, _ScanIter
2
3
4class SetCommandsMixin:
5    """Set commands mixin.
6
7    For commands details see: http://redis.io/commands#set
8    """
9
10    def sadd(self, key, member, *members):
11        """Add one or more members to a set."""
12        return self.execute(b'SADD', key, member, *members)
13
14    def scard(self, key):
15        """Get the number of members in a set."""
16        return self.execute(b'SCARD', key)
17
18    def sdiff(self, key, *keys):
19        """Subtract multiple sets."""
20        return self.execute(b'SDIFF', key, *keys)
21
22    def sdiffstore(self, destkey, key, *keys):
23        """Subtract multiple sets and store the resulting set in a key."""
24        return self.execute(b'SDIFFSTORE', destkey, key, *keys)
25
26    def sinter(self, key, *keys):
27        """Intersect multiple sets."""
28        return self.execute(b'SINTER', key, *keys)
29
30    def sinterstore(self, destkey, key, *keys):
31        """Intersect multiple sets and store the resulting set in a key."""
32        return self.execute(b'SINTERSTORE', destkey, key, *keys)
33
34    def sismember(self, key, member):
35        """Determine if a given value is a member of a set."""
36        return self.execute(b'SISMEMBER', key, member)
37
38    def smembers(self, key, *, encoding=_NOTSET):
39        """Get all the members in a set."""
40        return self.execute(b'SMEMBERS', key, encoding=encoding)
41
42    def smove(self, sourcekey, destkey, member):
43        """Move a member from one set to another."""
44        return self.execute(b'SMOVE', sourcekey, destkey, member)
45
46    def spop(self, key, count=None, *, encoding=_NOTSET):
47        """Remove and return one or multiple random members from a set."""
48        args = [key]
49        if count is not None:
50            args.append(count)
51        return self.execute(b'SPOP', *args, encoding=encoding)
52
53    def srandmember(self, key, count=None, *, encoding=_NOTSET):
54        """Get one or multiple random members from a set."""
55        args = [key]
56        count is not None and args.append(count)
57        return self.execute(b'SRANDMEMBER', *args, encoding=encoding)
58
59    def srem(self, key, member, *members):
60        """Remove one or more members from a set."""
61        return self.execute(b'SREM', key, member, *members)
62
63    def sunion(self, key, *keys):
64        """Add multiple sets."""
65        return self.execute(b'SUNION', key, *keys)
66
67    def sunionstore(self, destkey, key, *keys):
68        """Add multiple sets and store the resulting set in a key."""
69        return self.execute(b'SUNIONSTORE', destkey, key, *keys)
70
71    def sscan(self, key, cursor=0, match=None, count=None):
72        """Incrementally iterate Set elements."""
73        tokens = [key, cursor]
74        match is not None and tokens.extend([b'MATCH', match])
75        count is not None and tokens.extend([b'COUNT', count])
76        fut = self.execute(b'SSCAN', *tokens)
77        return wait_convert(fut, lambda obj: (int(obj[0]), obj[1]))
78
79    def isscan(self, key, *, match=None, count=None):
80        """Incrementally iterate set elements using async for.
81
82        Usage example:
83
84        >>> async for val in redis.isscan(key, match='something*'):
85        ...     print('Matched:', val)
86
87        """
88        return _ScanIter(lambda cur: self.sscan(key, cur,
89                                                match=match,
90                                                count=count))
91