Skip to content

Commit

Permalink
Improve typing for permutations and gnitb
Browse files Browse the repository at this point in the history
  • Loading branch information
inducer committed Jan 12, 2025
1 parent 2c22d9a commit 72732f4
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions pytools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1385,7 +1385,7 @@ def indices_in_shape(shape):


def generate_nonnegative_integer_tuples_below(
n: tuple[int, ...] | int, length: int | None = None, least: int = 0
n: Sequence[int] | int, length: int | None = None, least: int = 0
) -> Iterator[tuple[int, ...]]:
"""n may be a sequence, in which case length must be None."""
if length is None:
Expand Down Expand Up @@ -1477,7 +1477,10 @@ def generate_all_integer_tuples_below(
n, length, least_abs))


class _ConcatenableSequence(Protocol):
T_co = TypeVar("T_co", covariant=True)


class _ConcatenableSequence(Generic[T_co], Protocol):
"""
A protocol that supports the following:
Expand All @@ -1494,10 +1497,13 @@ def __add__(self, other: Self) -> Self:
def __len__(self) -> int:
...

def __iter__(self) -> Iterator[T_co]:
...


def generate_permutations(
original: _ConcatenableSequence
) -> Iterator[_ConcatenableSequence]:
original: _ConcatenableSequence[T]
) -> Iterator[_ConcatenableSequence[T]]:
"""Generate all permutations of the list *original*.
Nicked from http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/252178
Expand All @@ -1512,8 +1518,8 @@ def generate_permutations(


def generate_unique_permutations(
original: _ConcatenableSequence
) -> Iterator[_ConcatenableSequence]:
original: _ConcatenableSequence[T]
) -> Iterator[_ConcatenableSequence[T]]:
"""Generate all unique permutations of the list *original*.
Note that, unlike for :func:`generate_permutations`, *original* must be a
Expand Down

0 comments on commit 72732f4

Please sign in to comment.