server.py

#!/usr/bin/env python3
"""
02: RESP parser — Teach the server to speak Redis.

RESP (REdis Serialization Protocol) is the wire format real Redis uses.
Once the server can parse RESP, you can connect with `redis-cli`, not just nc.

Commands from a client always arrive as an array of bulk strings:
    *3\\r\\n$3\\r\\nSET\\r\\n$3\\r\\nfoo\\r\\n$5\\r\\nhello\\r\\n
    -> ["SET", "foo", "hello"]

This module parses incoming RESP and replies with a Bulk String containing
the parsed command. Step 3 will dispatch the command to a real handler.

Usage:
    uv run python 02-resp-parser/server.py
    redis-cli -p 6380 PING
    redis-cli -p 6380 SET foo hello
"""
from __future__ import annotations

import socket

HOST = "127.0.0.1"
PORT = 6380
CRLF = b"\r\n"


# --- RESP parser ---

class ProtocolError(ValueError):
    pass


def parse_resp(buf: bytes) -> tuple[object, int] | None:
    """
    Parse a single RESP frame from buf. Return (value, bytes_consumed) on
    success or None if the buffer doesn't contain a full frame yet.

    Supported types: Simple String (+), Error (-), Integer (:), Bulk String ($),
    Array (*).
    """
    if not buf:
        return None
    type_byte = buf[0:1]

    # Find the first CRLF that ends the header line
    end = buf.find(CRLF)
    if end < 0:
        return None
    header = buf[1:end].decode("ascii")
    after_header = end + 2

    if type_byte == b"+":  # Simple String
        return header, after_header

    if type_byte == b"-":  # Error
        return Exception(header), after_header

    if type_byte == b":":  # Integer
        return int(header), after_header

    if type_byte == b"$":  # Bulk String
        length = int(header)
        if length == -1:
            return None, after_header  # NULL bulk
        end_payload = after_header + length
        if len(buf) < end_payload + 2:  # +2 for trailing CRLF
            return None  # incomplete
        return buf[after_header:end_payload].decode("utf-8"), end_payload + 2

    if type_byte == b"*":  # Array
        count = int(header)
        if count == -1:
            return None, after_header  # NULL array
        items = []
        cursor = after_header
        for _ in range(count):
            result = parse_resp(buf[cursor:])
            if result is None:
                return None  # incomplete
            value, consumed = result
            items.append(value)
            cursor += consumed
        return items, cursor

    raise ProtocolError(f"Unknown RESP type byte: {type_byte!r}")


# --- RESP encoder ---

def encode_simple(s: str) -> bytes:
    return b"+" + s.encode("ascii") + CRLF


def encode_error(s: str) -> bytes:
    return b"-" + s.encode("ascii") + CRLF


def encode_bulk(s: str | None) -> bytes:
    if s is None:
        return b"$-1\r\n"
    data = s.encode("utf-8")
    return b"$" + str(len(data)).encode("ascii") + CRLF + data + CRLF


def encode_integer(n: int) -> bytes:
    return b":" + str(n).encode("ascii") + CRLF


# --- Server ---

def handle_client(client: socket.socket, addr: tuple[str, int]) -> None:
    print(f"  [conn] {addr} connected")
    buf = b""
    with client:
        while True:
            chunk = client.recv(4096)
            if not chunk:
                break
            buf += chunk
            # Parse as many complete frames as possible
            while True:
                result = parse_resp(buf)
                if result is None:
                    break  # need more bytes
                command, consumed = result
                buf = buf[consumed:]
                print(f"  [recv] {addr}: {command!r}")
                reply = handle_command(command)
                client.sendall(reply)
    print(f"  [conn] {addr} disconnected")


def handle_command(command: object) -> bytes:
    """Step 2 just echoes the parsed command. Step 3 will dispatch for real."""
    if isinstance(command, list) and command:
        name = command[0]
        if isinstance(name, str) and name.upper() == "PING":
            return encode_simple("PONG")
        return encode_bulk(f"{command}")
    return encode_error("ERR expected array")


def main() -> None:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind((HOST, PORT))
        server.listen(1)
        print(f"RESP server listening on {HOST}:{PORT}")
        while True:
            client, addr = server.accept()
            handle_client(client, addr)


if __name__ == "__main__":
    main()