✨ support asyncio

memcache/__init__.py

@@ -1,6 +1,8 @@
+from .async_memcache import AsyncMemcache
 from .memcache import Memcache, MetaCommand, MetaResult, MemcacheError
 
 __all__ = [
+    "AsyncMemcache",
     "Memcache",
     "MemcacheError",
     "MetaResult",

memcache/async_memcache.py

@@ -0,0 +1,98 @@
+import asyncio
+from typing import List, Tuple, Union, Optional
+
+from .memcache import Addr, MemcacheError, MetaCommand, MetaResult
+
+
+class AsyncConnection:
+    def __init__(self, addr: Union[Tuple[str, int]]):
+        self._addr = addr
+        self._connected = False
+
+    async def _connect(self) -> None:
+        self.reader, self.writer = await asyncio.open_connection(
+            self._addr[0], self._addr[1]
+        )
+        self._connected = True
+
+    async def flush_all(self) -> None:
+        if not self._connected:
+            await self._connect()
+
+        self.writer.write(b"flush_all\r\n")
+        await self.writer.drain()
+        response = await self.reader.readline()
+        if response != b"OK\r\n":
+            raise MemcacheError(response.removesuffix(b"\r\n"))
+
+    async def execute_meta_command(self, command: MetaCommand) -> MetaResult:
+        if not self._connected:
+            await self._connect()
+
+        header = b" ".join([command.cm, command.key] + command.flags + [b"\r\n"])
+        self.writer.write(header)
+        if command.value:
+            self.writer.write(command.value + b"\r\n")
+        await self.writer.drain()
+        return await self._receive_meta_result()
+
+    async def _receive_meta_result(self) -> MetaResult:
+        header = await self.reader.readline()
+        parts = header.split()
+        rc = parts[0]
+        flags = parts[1:]
+        value = None
+        if rc == b"VA":
+            size = int(parts[1])
+            flags = parts[2:]
+            value = await self.reader.read(size)
+            await self.reader.read(2)  # read the "\r\n"
+        return MetaResult(rc=rc, flags=flags, value=value)
+
+    async def set(
+        self, key: Union[bytes, str], value: bytes, expire: Optional[int] = None
+    ) -> None:
+        flags = [b"S%d" % len(value)]
+        if expire:
+            flags.append(b"T%d" % expire)
+
+        command = MetaCommand(cm=b"ms", key=key, flags=flags, value=value)
+        await self.execute_meta_command(command)
+
+    async def get(self, key: Union[bytes, str]) -> Optional[bytes]:
+        command = MetaCommand(cm=b"mg", key=key, flags=[b"v"], value=None)
+        return (await self.execute_meta_command(command)).value
+
+    async def delete(self, key: Union[bytes, str]) -> None:
+        command = MetaCommand(cm=b"md", key=key, flags=[], value=None)
+        await self.execute_meta_command(command)
+
+
+class AsyncMemcache:
+    def __init__(self, addr: Union[Addr, List[Addr]] = None):
+        addr = addr or ("localhost", 11211)
+        if isinstance(addr, list):
+            self.connections = [AsyncConnection(x) for x in addr]
+        else:
+            self.connections = [AsyncConnection(addr)]
+
+    def _get_connection(self, key) -> AsyncConnection:
+        return self.connections[hash(key) % len(self.connections)]
+
+    async def execute_meta_command(self, command: MetaCommand) -> MetaResult:
+        return await self._get_connection(command.key).execute_meta_command(command)
+
+    async def flush_all(self) -> None:
+        for connection in self.connections:
+            await connection.flush_all()
+
+    async def set(
+        self, key: Union[bytes, str], value: bytes, *, expire: Optional[int] = None
+    ) -> None:
+        return await self._get_connection(key).set(key, value, expire=expire)
+
+    async def get(self, key: Union[bytes, str]) -> Optional[bytes]:
+        return await self._get_connection(key).get(key)
+
+    async def delete(self, key: Union[bytes, str]) -> None:
+        return await self._get_connection(key).delete(key)

pyproject.toml

@@ -12,6 +12,7 @@ python = "^3.6"
 flake8 = "^3.8.4"
 black = "^20.8b1"
 pytest = "^6.2.1"
+pytest-asyncio = "^0.15.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/test_async_client.py

@@ -0,0 +1,44 @@
+import asyncio
+import time
+
+import pytest
+
+import memcache
+
+
+@pytest.fixture()
+def client():
+    return memcache.AsyncMemcache(("localhost", 11211))
+
+
+@pytest.mark.asyncio
+async def test_execute_command(client):
+    command = memcache.MetaCommand(cm=b"ms", key=b"foo", flags=[b"S3"], value=b"bar")
+    result = await client.execute_meta_command(command)
+    assert result.rc == b"OK"
+
+    command = memcache.MetaCommand(cm=b"mg", key=b"foo", flags=[b"v"], value=None)
+    result = await client.execute_meta_command(command)
+    assert result.rc == b"VA"
+    assert result.value == b"bar"
+
+
+@pytest.mark.asyncio
+async def test_set_get(client):
+    await client.set("foofoo", b"barbar", expire=1)
+    assert await client.get(b"foofoo") == b"barbar"
+    await asyncio.sleep(1.1)
+    assert await client.get(b"foofoo") is None
+
+
+@pytest.mark.asyncio
+async def test_delete(client):
+    await client.set(b"to_be_deleted", b"xxxx")
+    await client.delete(b"to_be_deleted")
+    assert await client.get(b"to_be_deleted") is None
+
+
+@pytest.mark.asyncio
+async def test_flush_all(client):
+    await client.flush_all()
+    await client.flush_all()