vi.util

 1from __future__ import annotations
 2
 3import random
 4from typing import TYPE_CHECKING
 5
 6from pygame.math import Vector2
 7
 8
 9if TYPE_CHECKING:
10    from collections.abc import Iterator
11
12    from pygame.rect import Rect
13
14
15__all__ = [
16    "count",
17    "first",
18    "probability",
19    "random_angle",
20    "random_pos",
21    "round_pos",
22]
23
24
25def probability(threshold: float, prng: random.Random | None = None) -> bool:
26    """Randomly retrieve True or False depending on the given probability.
27
28    The probability should be between 0 and 1.
29    If you give a probability equal or higher than 1, this function will always return True.
30    Likewise, if you give a probability equal or lower than 0, this function will always return False.
31    """
32    get_random = prng.random if prng else random.random
33    return threshold > get_random()
34
35
36def round_pos(pos: Vector2) -> tuple[int, int]:
37    """Round the x and y coordinates of a vector to two integers respectively."""
38    return round(pos.x), round(pos.y)
39
40
41def random_angle(length: float, prng: random.Random | None = None) -> Vector2:
42    """Retrieve a randomly-angled vector with a given length."""
43    uniform = prng.uniform if prng else random.uniform
44
45    vec = Vector2(length, 0)
46    vec.rotate_ip(uniform(0, 360))
47    return vec
48
49
50def random_pos(area: Rect, prng: random.Random | None = None) -> Vector2:
51    """Retrieve a random position within an area."""
52    uniform = prng.uniform if prng else random.uniform
53
54    x = uniform(area.left, area.right)
55    y = uniform(area.top, area.bottom)
56    return Vector2(x, y)
57
58
59def first[T](iterator: Iterator[T]) -> T | None:
60    """Returns the first element in an iterator.
61
62    Returns None if the iterator contains no elements.
63    """
64    return next(iterator, None)
65
66
67def count[T](iterator: Iterator[T]) -> int:
68    """Count the number of elements in an iterator.
69
70    An alternative way to count the number of elements in an iterator is to collect all the elements
71    in a list first and then retrieve its length. However, a memory allocation will be used, which
72    can hurt performance.
73
74    ```python
75    class MyAgent(Agent):
76        def update(self) -> None:
77            count = len(list(self.in_proximity_accuracy()))
78    ```
79    """
80    return sum(1 for _ in iterator)