from collections import deque
from typing import *
import asyncio
import collections
import itertools
import sys
import time
from typing import Any
import hiredis
import uvloop
expiration = collections.defaultdict(lambda: float("inf")) # type: Dict[bytes, float]
dictionary = {} # type: Dict[bytes, Any]
class RedisProtocol(asyncio.Protocol):
def __init__(self):
self.dictionary = dictionary
self.response = collections.deque()
self.parser = hiredis.Reader()
self.transport = None # type: asyncio.transports.Transport
self.commands = {
b"COMMAND": self.command,
b"SET": self.set,
b"GET": self.get,
b"PING": self.ping,
b"INCR": self.incr,
b"LPUSH": self.lpush,
b"RPUSH": self.rpush,
b"LPOP": self.lpop,
b"RPOP": self.rpop,
b"SADD": self.sadd,
b"HSET": self.hset,
b"SPOP": self.spop,
b"LRANGE": self.lrange,
b"MSET": self.mset,
}
def connection_made(self, transport: asyncio.transports.Transport):
self.transport = transport
def data_received(self, data: bytes):
self.parser.feed(data)
while 1:
req = self.parser.gets()
if req is False:
break
else:
self.response.append(self.commands[req[0].upper()](*req[1:]))
self.transport.writelines(self.response)
self.response.clear()
def command(self):
# Far from being a complete implementation of the `COMMAND` command of
# Redis, yet sufficient for us to start using redis-cli.
return b"+OK\r\n"
def set(self, *args) -> bytes:
# Defaults
key = args[0]
value = args[1]
expires_at = None
cond = b""
largs = len(args)
if largs == 3:
# SET key value [NX|XX]
cond = args[2]
elif largs >= 4:
# SET key value [EX seconds | PX milliseconds] [NX|XX]
try:
if args[2] == b"EX":
duration = int(args[3])
elif args[2] == b"PX":
duration = int(args[3]) / 1000
else:
return b"-ERR syntax error\r\n"
except ValueError:
return b"-value is not an integer or out of range\r\n"
if duration <= 0:
return b"-ERR invalid expire time in set\r\n"
expires_at = time.monotonic() + duration
if largs == 5:
cond = args[4]
if cond == b"":
pass
elif cond == b"NX":
if key in self.dictionary:
return b"$-1\r\n"
elif cond == b"XX":
if key not in self.dictionary:
return b"$-1\r\n"
else:
return b"-ERR syntax error\r\n"
if expires_at:
expiration[key] = expires_at
self.dictionary[key] = value
return b"+OK\r\n"
def get(self, key: bytes) -> bytes:
if key not in self.dictionary:
return b"$-1\r\n"
if key in expiration and expiration[key] < time.monotonic():
del self.dictionary[key]
del expiration[key]
return b"$-1\r\n"
else:
value = self.dictionary[key]
return b"$%d\r\n%s\r\n" % (len(value), value)
def ping(self, message=b"PONG"):
return b"$%d\r\n%s\r\n" % (len(message), message)
def incr(self, key):
value = self.dictionary.get(key, 0)
if type(value) is str:
try:
value = int(value)
except ValueError:
return b"-value is not an integer or out of range\r\n"
value += 1
self.dictionary[key] = str(value)
return b":%d\r\n" % (value,)
def lpush(self, key, *values):
deque = self.dictionary.get(key, collections.deque())
deque.extendleft(values)
self.dictionary[key] = deque
return b":%d\r\n" % (len(deque),)
def rpush(self, key, *values):
deque = self.dictionary.get(key, collections.deque())
deque.extend(values)
self.dictionary[key] = deque
return b":%d\r\n" % (len(deque),)
def lpop(self, key):
try:
deque = self.dictionary[key] # type: collections.deque
except KeyError:
return b"$-1\r\n"
value = deque.popleft()
return b"$%d\r\n%s\r\n" % (len(value), value)
def rpop(self, key):
try:
deque = self.dictionary[key] # type: collections.deque
except KeyError:
return b"$-1\r\n"
value = deque.pop()
return b"$%d\r\n%s\r\n" % (len(value), value)
def sadd(self, key, *members):
set_ = self.dictionary.get(key, set())
prev_size = len(set_)
for member in members:
set_.add(member)
self.dictionary[key] = set_
return b":%d\r\n" % (len(set_) - prev_size,)
def hset(self, key, field, value):
hash_ = self.dictionary.get(key, {})
ret = int(field in hash_)
hash_[field] = value
self.dictionary[key] = hash_
return b":%d\r\n" % (ret,)
def spop(self, key): # TODO add `count`
try:
set_ = self.dictionary[key] # type: set
elem = set_.pop()
except KeyError:
return b"$-1\r\n"
return b"$%d\r\n%s\r\n" % (len(elem), elem)
def lrange(self, key, start, stop):
start = int(start)
stop = int(stop)
try:
deque = self.dictionary[key] # type: collections.deque
except KeyError:
return b"$-1\r\n"
l = itertools.islice(deque, start, stop)
return b"*%d\r\n%s" % (stop - start, b"".join(b"$%d\r\n%s\r\n" % (len(e), e) for e in l))
def mset(self, *args):
for i in range(0, len(args), 2):
key = args[i]
value = args[i + 1]
self.dictionary[key] = value
return b"+OK\r\n"
def main() -> int:
print("Hello, World!")
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
loop = asyncio.get_event_loop()
# Each client connection will create a new protocol instance
coro = loop.create_server(RedisProtocol, "127.0.0.1", 7878)
server = loop.run_until_complete(coro)
# Serve requests until Ctrl+C is pressed
print('Serving on {}'.format(server.sockets[0].getsockname()))
try:
loop.run_forever()
except KeyboardInterrupt:
pass
# Close the server
server.close()
loop.run_until_complete(server.wait_closed())
loop.close()
return 0
if __name__ == "__main__":
sys.exit(main())
