Coverage for o2/util/helper.py: 91%
67 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-16 11:18 +0000
1import functools
2import random
3import re
4import string
5from collections.abc import Generator
6from typing import TYPE_CHECKING, Any, Callable, Concatenate, Optional, ParamSpec, TypeVar, Union
8import numpy as np
9import xxhash
10from sympy import Symbol, lambdify
12from o2.models.settings import ActionVariationSelection
14if TYPE_CHECKING:
15 from o2.store import Store
17CLONE_REGEX = re.compile(r"^(.*)_clone_[a-z0-9]{8}(?:timetable)?$")
20def random_string(length: int = 8) -> str:
21 """Generate a random alphanumeric string of length `length`."""
22 return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))
25def name_is_clone_of(potential_clone_name: str, resource_id: str) -> bool:
26 """Check if the name or id is a clone of a resource id."""
27 match = CLONE_REGEX.match(potential_clone_name)
28 return match is not None and (
29 match.group(1) == resource_id
30 or match.group(1) + "timetable" == resource_id
31 or match.group(1) == resource_id + "timetable"
32 )
35T = TypeVar("T")
38def safe_list_index(lst: list[T], item: T) -> Optional[int]:
39 """Return the index of the item in the list, or None if it is not present."""
40 try:
41 return lst.index(item)
42 except ValueError:
43 return None
46def hash_int(s: object) -> int:
47 """Create int hash based on the string representation of the object."""
48 return xxhash.xxh3_64_intdigest(str(s))
51def hash_string(s: object) -> str:
52 """Create string hash based on the string representation of the object."""
53 return xxhash.xxh3_64_hexdigest(str(s)).zfill(16)
56def hex_id(id: Union[int, np.int64]) -> str:
57 """Convert an item id to a hex id."""
58 # If the item id is negative, then it "overflowed" to a signed int,
59 # so we need to convert it to an unsigned int.
60 # (rTree used C types in the backend, so int don't "just" scale)
61 if isinstance(id, np.number):
62 id = id.item()
63 if id < 0:
64 id &= 0xFFFFFFFFFFFFFFFF
65 return f"{id:x}".zfill(16)
68@functools.lru_cache(maxsize=100)
69def cached_lambdify(expr: str) -> Callable[[float], float]:
70 """Lambdify an expression and cache the result."""
71 return functools.lru_cache(maxsize=100)(lambdify(Symbol("size"), expr))
74def lambdify_dict(d: dict[str, str]) -> dict[str, Callable[[float], float]]:
75 """Convert all lambdas in the dictionary to functions."""
76 return {k: cached_lambdify(v) for k, v in d.items()}
79P = ParamSpec("P")
80R = TypeVar("R")
83def with_signature_from(
84 f: Callable[Concatenate[Any, P], R], /
85) -> Callable[[Callable[Concatenate[Any, P], R]], Callable[Concatenate[Any, P], R]]:
86 """Copy the signature from one function to another.
88 Allows creating a function that has the same signature as another function,
89 which is useful for creating wrappers while preserving type hinting.
90 """
91 return lambda _: _
94P = TypeVar("P")
97def select_variant(
98 store: "Store",
99 options: list[P],
100 inner: Optional[bool] = False,
101 ordered: Optional[bool] = False,
102) -> list[P]:
103 """Pick a single or multiple elements from the list.
105 This depends on the action_variation_selection setting.
106 The inner parameter can be used to signal, that this is an inner loop,
107 and so we should at max pick 1 element.
109 If inner is True and the action_variation_selection is
110 FIRST_MAX_VARIANTS_PER_ACTION_IN_ORDER or RANDOM_MAX_VARIANTS_PER_ACTION,
111 we will pick at max 1 element. This is used to limit the number of elements
112 we need to consider in the inner loop.
114 The ordered parameter can be used to signal, that the options are already ordered,
115 so random selection is not needed.
116 """
117 if not options:
118 return []
120 action_variation_selection = store.settings.action_variation_selection
122 if ordered:
123 action_variation_selection = action_variation_selection.ordered
124 if inner and not store.settings.override_action_variation_selection_for_inner_loop:
125 action_variation_selection = action_variation_selection.inner
127 if action_variation_selection == ActionVariationSelection.SINGLE_RANDOM:
128 return random.sample(options, 1)
129 elif action_variation_selection == ActionVariationSelection.ALL_RANDOM:
130 return random.sample(options, len(options))
131 elif action_variation_selection == ActionVariationSelection.FIRST_IN_ORDER:
132 return options[:1]
133 elif action_variation_selection == ActionVariationSelection.FIRST_MAX_VARIANTS_PER_ACTION_IN_ORDER:
134 assert isinstance(store.settings.max_variants_per_action, int)
135 return options[: store.settings.max_variants_per_action]
136 elif action_variation_selection == ActionVariationSelection.RANDOM_MAX_VARIANTS_PER_ACTION:
137 assert isinstance(store.settings.max_variants_per_action, int)
138 return random.sample(options, min(store.settings.max_variants_per_action, len(options)))
139 elif action_variation_selection == ActionVariationSelection.ALL_IN_ORDER:
140 return options
141 else:
142 raise ValueError(f"Invalid action variation selection: {action_variation_selection}")
145def select_variants(
146 store: "Store",
147 options: list[P],
148 inner: Optional[bool] = False,
149 ordered: Optional[bool] = False,
150) -> Generator[P, None, None]:
151 """Create a generator (to be used in a for loop) that yields action variations.
153 In accordance with the action_variation_selection setting.
155 The inner parameter can be used to signal, that this is an inner loop,
156 and (if action_variation_selection is FIRST_MAX_VARIANTS_PER_ACTION_IN_ORDER or
157 RANDOM_MAX_VARIANTS_PER_ACTION) we should pick at max 1 element.
159 The ordered parameter can be used to signal, that the options are already ordered,
160 so random selection is not needed.
161 """
162 yield from select_variant(store, options, inner, ordered)