Type hinting to return the same type of Sized iterable as the input?

Question

The following code works, but depending on how I fix the type hinting, either PyCharm or mypy complains about it. I've tried Sized, Iterable, and Collection for the type of S.

T = TypeVar("T")
S = TypeVar("S", bound=Collection[T])


def limit(i: S, n: int) -> S:
    """
    Limits the size of the input iterable to n.  Truncation is randomly chosen.
    """
    if len(i) <= n:
        return i
    return type(i)(random.sample(list(i), n))

What I want is something like random.sample with the following two properties:

1. If n > len(i), return i instead of throwing an error
2. The type of Collection of the output should be identical to the input. list[int | str] in = list[int | str] out; set[float] in = set[float] out.

If there's already something like this in the standard library, I'll use that.

Answer 1

You don't need T here. Just by declaring a bound of Collection[Any], S can take on any "actual" collection form:

_{(playgrounds: Mypy, Pyright)}

def limit[S: Collection[Any]](i: S) -> S: ...

reveal_type(limit([1, 2, 3]))             # list[int]
reveal_type(limit({1, 2, 3}))             # set[int]
reveal_type(limit(frozenset({1, 2, 3})))  # frozenset[int]

However, Collection doesn't guarantee anything about its constructor, so invoking type(i)(list(...)) is an error. Replacing Collection[Any] with list[Any] | set[Any] | frozenset[Any] should help eliminating this error:

def limit[S: list[Any] | set[Any] | frozenset[Any]](i: S, n: int) -> S: ...

In reality, it doesn't:

_{(playgrounds: Mypy, Pyright)}

reveal_type(type(i))  # type[S]

# pyright => fine
# mypy => error: Incompatible return value type (got "list[Any] | set[Any] | frozenset[Any]", expected "S")
return type(i)(random.sample(list(i), n))

This is either a bug or limitation of Mypy, because a type[S] should produce a S instead of its bound on construction. For now, ignore this with a # type: ignore.

In conclusion:

_{(playgrounds: Mypy, Pyright)}

from typing import Any
import random

def limit[S: list[Any] | set[Any] | frozenset[Any]](i: S, n: int) -> S:
    if len(i) <= n:
        return i
    
    return type(i)(random.sample(list(i), n))  # type: ignore[return-value]

reveal_type(limit([42, '', 3.14], 0))             # list[int | str | float] / list[object]
reveal_type(limit({42, '', 3.14}, 0))             # set[int | str | float] / set[object]
reveal_type(limit(frozenset({42, '', 3.14}), 0))  # frozenset[Any] / frozenset[object]

type SomeComplexUnion = list[int | str | float] | set[int | str | float]
v = cast(SomeComplexUnion, [42, '', 3.14])

reveal_type(v)                                    # SomeComplexUnion
reveal_type(limit(v, 0))                          # SomeComplexUnion

Pre-PEP-695 syntax:

S = TypeVar('S', bound = list[Any] | set[Any] | frozenset[Any])

def limit(i: S, n: int) -> S: ...