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:
n > len(i)
, return i
instead of throwing an errorlist[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.
You don't need T
here. Just by declaring a bound of Collection[Any]
, S
can take on any "actual" collection form:
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:
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:
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: ...