overload.py

from typing import *

class RawStream:
    pass

class Stream:
    pass

class Other:
    pass

_T = TypeVar('_T')

Factory = Callable[[RawStream], Awaitable[_T]]

async def default(raw: RawStream) -> Stream:
    return Stream()

async def none(raw: RawStream) -> RawStream:
    return raw

async def other(raw: RawStream) -> Other:
    return Other()


@overload
async def connect() -> Stream: ...

@overload
async def connect(factory: Factory[_T]) -> _T: ...

async def connect(factory=default):  # type: ignore
    raw = RawStream()
    return await factory(raw)


async def test() -> None:
    s: Stream = await connect()
    r: RawStream = await connect(none)
    o = await connect(other)
    reveal_type(o)
    await connect(None)