stricter typing for reader._parser for https://github.com/lemon24/reader/issues/271

strict_parser.py

"""
stricter typing 
for https://github.com/lemon24/reader/blob/master/src/reader/_parser.py
for https://github.com/lemon24/reader/issues/271

"""
from typing import *
import io
from dataclasses import dataclass
from random import choice

T = TypeVar('T')
T_co = TypeVar('T_co', covariant=True)
T_cv = TypeVar('T_cv', contravariant=True)


@dataclass
class RetrieveResult(Generic[T_co]):
    payload: T_co
    mime_type: Optional[str] = None

class RetrieverType(Protocol[T_co]):
    def __call__(self) -> RetrieveResult[T_co]: ...

class ParserType(Protocol[T_cv]):
    def __call__(self, payload: T_cv) -> str: ...
    

# in general, there are many types of retrievers that produce
# RetrieveResult[IO[bytes]], and many types of parsers that consume it;
# the parser is selected based on the result mime_type

class FileRetriever:
    def __call__(self) -> RetrieveResult[IO[bytes]]:
        return RetrieveResult(io.BytesIO(b'from-file'), choice(['xml', 'other']))

class HTTPRetriever:
    def __call__(self) -> RetrieveResult[IO[bytes]]:
        return RetrieveResult(io.BytesIO(b'from-http'), choice(['xml', 'other']))

class FeedParser:
    def __call__(self, payload: IO[bytes]) -> str:
        return 'feed-' + payload.read().decode()

# ... but for a custom use case, I have a tightly coupled 
# retriever/parser pair for which the payload has a specific type;
# I would like mypy to check that type, if possible

class CustomRetriever:
    def __call__(self) -> RetrieveResult[bytes]:
        return RetrieveResult(b'from-custom', 'x.custom')
    
class CustomParser:
    def __call__(self, payload: bytes) -> str:
        return 'custom-' + payload.decode()


# mime_type -> parser; new parsers may be added by the user
PARSERS: Dict[str, ParserType[Any]] = {
    #'fallback': FeedParser(),
    #'xml': FeedParser(),
    #'x.custom': CustomParser(),
    
    # AttributeError: '_io.BytesIO' object has no attribute 'decode'
    # oops, this is a bug, how do I get mypy to complain?
    'fallback': FeedParser(),
    'x.custom': FeedParser(),
    'xml': CustomParser(),
}

def get_parser(result: RetrieveResult[T]) -> ParserType[T]:
    if result.mime_type not in PARSERS:
        return PARSERS['fallback']
    rv = PARSERS[result.mime_type]
    #reveal_type(rv)
    return rv


def file_fn() -> None:
    result = FileRetriever()()
    rv = get_parser(result)(result.payload)
    print(repr(rv))
    
def http_fn() -> None:
    result = HTTPRetriever()()
    rv = get_parser(result)(result.payload)
    print(repr(rv))
    
def custom_fn() -> None:
    result = CustomRetriever()()
    rv = get_parser(result)(result.payload)
    print(repr(rv))
    
file_fn()
http_fn()
custom_fn()