1import copy
2import itertools
3import operator
4from collections.abc import Callable
5from functools import total_ordering, wraps
6from typing import Any, TypeVar
7
8T = TypeVar("T")
9
10
11class classproperty:
12 """
13 Decorator that converts a method with a single cls argument into a property
14 that can be accessed directly from the class.
15 """
16
17 def __init__(self, method: Callable[[type], Any] | None = None) -> None:
18 self.fget = method
19
20 def __get__(self, instance: Any, cls: type | None = None) -> Any:
21 assert cls is not None, "cls must be provided for classproperty"
22 assert self.fget is not None, "fget must be set before accessing classproperty"
23 return self.fget(cls)
24
25 def getter(self, method: Callable[[type], Any]) -> classproperty:
26 self.fget = method
27 return self
28
29
30class Promise:
31 """
32 Base class for the proxy class created in the closure of the lazy function.
33 It's used to recognize promises in code.
34 """
35
36
37def lazy(func: Callable[..., Any], *resultclasses: type) -> Callable[..., Any]:
38 """
39 Turn any callable into a lazy evaluated callable. result classes or types
40 is required -- at least one is needed so that the automatic forcing of
41 the lazy evaluation code is triggered. Results are not memoized; the
42 function is evaluated on every access.
43 """
44
45 @total_ordering
46 class __proxy__(Promise):
47 """
48 Encapsulate a function call and act as a proxy for methods that are
49 called on the result of that function. The function is not evaluated
50 until one of the methods on the result is called.
51 """
52
53 __prepared = False
54
55 def __init__(self, args: tuple[Any, ...], kw: dict[str, Any]) -> None:
56 self.__args = args
57 self.__kw = kw
58 if not self.__prepared:
59 self.__prepare_class__()
60 self.__class__.__prepared = True
61
62 def __reduce__(self) -> tuple[Callable[..., Any], tuple[Any, ...]]:
63 return (
64 _lazy_proxy_unpickle,
65 (func, self.__args, self.__kw) + resultclasses,
66 )
67
68 def __repr__(self) -> str:
69 return repr(self.__cast())
70
71 @classmethod
72 def __prepare_class__(cls) -> None:
73 for resultclass in resultclasses:
74 for type_ in resultclass.mro():
75 for method_name in type_.__dict__:
76 # All __promise__ return the same wrapper method, they
77 # look up the correct implementation when called.
78 if hasattr(cls, method_name):
79 continue
80 meth = cls.__promise__(method_name)
81 setattr(cls, method_name, meth)
82 cls._delegate_bytes = bytes in resultclasses
83 cls._delegate_text = str in resultclasses
84 if cls._delegate_bytes and cls._delegate_text:
85 raise ValueError(
86 "Cannot call lazy() with both bytes and text return types."
87 )
88 if cls._delegate_text:
89 setattr(cls, "__str__", cls.__text_cast)
90 elif cls._delegate_bytes:
91 setattr(cls, "__bytes__", cls.__bytes_cast)
92
93 @classmethod
94 def __promise__(cls, method_name: str) -> Callable[..., Any]:
95 # Builds a wrapper around some magic method
96 def __wrapper__(self: Any, *args: Any, **kw: Any) -> Any:
97 # Automatically triggers the evaluation of a lazy value and
98 # applies the given magic method of the result type.
99 res = func(*self.__args, **self.__kw)
100 return getattr(res, method_name)(*args, **kw)
101
102 return __wrapper__
103
104 def __text_cast(self) -> str:
105 return func(*self.__args, **self.__kw)
106
107 def __bytes_cast(self) -> bytes:
108 return bytes(func(*self.__args, **self.__kw))
109
110 def __bytes_cast_encoded(self) -> bytes:
111 return func(*self.__args, **self.__kw).encode()
112
113 def __cast(self) -> Any:
114 if self._delegate_bytes:
115 return self.__bytes_cast()
116 elif self._delegate_text:
117 return self.__text_cast()
118 else:
119 return func(*self.__args, **self.__kw)
120
121 def __str__(self) -> str:
122 # object defines __str__(), so __prepare_class__() won't overload
123 # a __str__() method from the proxied class.
124 return str(self.__cast())
125
126 def __eq__(self, other: object) -> bool:
127 if isinstance(other, Promise):
128 # The mangled __cast exists on the runtime proxy subclass.
129 other = other.__cast() # ty: ignore[unresolved-attribute]
130 return self.__cast() == other
131
132 def __lt__(self, other: Any) -> bool:
133 if isinstance(other, Promise):
134 other = other.__cast()
135 return self.__cast() < other
136
137 def __hash__(self) -> int:
138 return hash(self.__cast())
139
140 def __mod__(self, rhs: Any) -> Any:
141 if self._delegate_text:
142 return str(self) % rhs
143 return self.__cast() % rhs
144
145 def __add__(self, other: Any) -> Any:
146 return self.__cast() + other
147
148 def __radd__(self, other: Any) -> Any:
149 return other + self.__cast()
150
151 def __deepcopy__(self, memo: dict[int, Any]) -> __proxy__:
152 # Instances of this class are effectively immutable. It's just a
153 # collection of functions. So we don't need to do anything
154 # complicated for copying.
155 memo[id(self)] = self
156 return self
157
158 @wraps(func)
159 def __wrapper__(*args: Any, **kw: Any) -> __proxy__:
160 # Creates the proxy object, instead of the actual value.
161 return __proxy__(args, kw)
162
163 return __wrapper__
164
165
166def _lazy_proxy_unpickle(
167 func: Callable[..., Any],
168 args: tuple[Any, ...],
169 kwargs: dict[str, Any],
170 *resultclasses: type,
171) -> Any:
172 return lazy(func, *resultclasses)(*args, **kwargs)
173
174
175def keep_lazy(
176 *resultclasses: type,
177) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
178 """
179 A decorator that allows a function to be called with one or more lazy
180 arguments. If none of the args are lazy, the function is evaluated
181 immediately, otherwise a __proxy__ is returned that will evaluate the
182 function when needed.
183 """
184 if not resultclasses:
185 raise TypeError("You must pass at least one argument to keep_lazy().")
186
187 def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
188 lazy_func = lazy(func, *resultclasses)
189
190 @wraps(func)
191 def wrapper(*args: Any, **kwargs: Any) -> Any:
192 if any(
193 isinstance(arg, Promise)
194 for arg in itertools.chain(args, kwargs.values())
195 ):
196 return lazy_func(*args, **kwargs)
197 return func(*args, **kwargs)
198
199 return wrapper
200
201 return decorator
202
203
204def keep_lazy_text(func: Callable[..., Any]) -> Callable[..., Any]:
205 """
206 A decorator for functions that accept lazy arguments and return text.
207 """
208 return keep_lazy(str)(func)
209
210
211empty = object()
212
213
214def new_method_proxy(func: Callable[..., Any]) -> Callable[..., Any]:
215 def inner(self: Any, *args: Any) -> Any:
216 if (_wrapped := self._wrapped) is empty:
217 self._setup()
218 _wrapped = self._wrapped
219 return func(_wrapped, *args)
220
221 inner._mask_wrapped = False # ty: ignore[unresolved-attribute]
222 return inner
223
224
225class LazyObject:
226 """
227 A wrapper for another class that can be used to delay instantiation of the
228 wrapped class.
229
230 By subclassing, you have the opportunity to intercept and alter the
231 instantiation. If you don't need to do that, use SimpleLazyObject.
232 """
233
234 # Avoid infinite recursion when tracing __init__ (#19456).
235 _wrapped = None
236
237 def __init__(self):
238 # Note: if a subclass overrides __init__(), it will likely need to
239 # override __copy__() and __deepcopy__() as well.
240 self._wrapped = empty
241
242 def __getattribute__(self, name: str) -> Any:
243 if name == "_wrapped":
244 # Avoid recursion when getting wrapped object.
245 return super().__getattribute__(name)
246 value = super().__getattribute__(name)
247 # If attribute is a proxy method, raise an AttributeError to call
248 # __getattr__() and use the wrapped object method.
249 if not getattr(value, "_mask_wrapped", True):
250 raise AttributeError
251 return value
252
253 __getattr__ = new_method_proxy(getattr)
254
255 def __setattr__(self, name: str, value: Any) -> None:
256 if name == "_wrapped":
257 # Assign to __dict__ to avoid infinite __setattr__ loops.
258 self.__dict__["_wrapped"] = value
259 else:
260 if self._wrapped is empty:
261 self._setup()
262 setattr(self._wrapped, name, value)
263
264 def __delattr__(self, name: str) -> None:
265 if name == "_wrapped":
266 raise TypeError("can't delete _wrapped.")
267 if self._wrapped is empty:
268 self._setup()
269 delattr(self._wrapped, name)
270
271 def _setup(self) -> None:
272 """
273 Must be implemented by subclasses to initialize the wrapped object.
274 """
275 raise NotImplementedError(
276 "subclasses of LazyObject must provide a _setup() method"
277 )
278
279 # Because we have messed with __class__ below, we confuse pickle as to what
280 # class we are pickling. We're going to have to initialize the wrapped
281 # object to successfully pickle it, so we might as well just pickle the
282 # wrapped object since they're supposed to act the same way.
283 #
284 # Unfortunately, if we try to simply act like the wrapped object, the ruse
285 # will break down when pickle gets our id(). Thus we end up with pickle
286 # thinking, in effect, that we are a distinct object from the wrapped
287 # object, but with the same __dict__. This can cause problems (see #25389).
288 #
289 # So instead, we define our own __reduce__ method and custom unpickler. We
290 # pickle the wrapped object as the unpickler's argument, so that pickle
291 # will pickle it normally, and then the unpickler simply returns its
292 # argument.
293 def __reduce__(self) -> tuple[Callable[[Any], Any], tuple[Any, ...]]:
294 if self._wrapped is empty:
295 self._setup()
296 return (unpickle_lazyobject, (self._wrapped,))
297
298 def __copy__(self) -> LazyObject | Any:
299 if self._wrapped is empty:
300 # If uninitialized, copy the wrapper. Use type(self), not
301 # self.__class__, because the latter is proxied.
302 return type(self)()
303 else:
304 # If initialized, return a copy of the wrapped object.
305 return copy.copy(self._wrapped)
306
307 def __deepcopy__(self, memo: dict[int, Any]) -> LazyObject | Any:
308 if self._wrapped is empty:
309 # We have to use type(self), not self.__class__, because the
310 # latter is proxied.
311 result = type(self)()
312 memo[id(self)] = result
313 return result
314 return copy.deepcopy(self._wrapped, memo)
315
316 __bytes__ = new_method_proxy(bytes)
317 __str__ = new_method_proxy(str)
318 __bool__ = new_method_proxy(bool)
319
320 # Introspection support
321 __dir__ = new_method_proxy(dir)
322
323 # Need to pretend to be the wrapped class, for the sake of objects that
324 # care about this (especially in equality tests)
325 __class__ = property(new_method_proxy(operator.attrgetter("__class__")))
326 __eq__ = new_method_proxy(operator.eq)
327 __lt__ = new_method_proxy(operator.lt)
328 __gt__ = new_method_proxy(operator.gt)
329 __ne__ = new_method_proxy(operator.ne)
330 __hash__ = new_method_proxy(hash)
331
332 # List/Tuple/Dictionary methods support
333 __getitem__ = new_method_proxy(operator.getitem)
334 __setitem__ = new_method_proxy(operator.setitem)
335 __delitem__ = new_method_proxy(operator.delitem)
336 __iter__ = new_method_proxy(iter)
337 __len__ = new_method_proxy(len)
338 __contains__ = new_method_proxy(operator.contains)
339
340
341def unpickle_lazyobject(wrapped: Any) -> Any:
342 """
343 Used to unpickle lazy objects. Just return its argument, which will be the
344 wrapped object.
345 """
346 return wrapped
347
348
349class SimpleLazyObject(LazyObject):
350 """
351 A lazy object initialized from any function.
352
353 Designed for compound objects of unknown type. For builtins or objects of
354 known type, use plain.utils.functional.lazy.
355 """
356
357 def __init__(self, func: Callable[[], Any]) -> None:
358 """
359 Pass in a callable that returns the object to be wrapped.
360
361 If copies are made of the resulting SimpleLazyObject, which can happen
362 in various circumstances within Plain, then you must ensure that the
363 callable can be safely run more than once and will return the same
364 value.
365 """
366 self.__dict__["_setupfunc"] = func
367 super().__init__()
368
369 def _setup(self) -> None:
370 self._wrapped = self._setupfunc()
371
372 # Return a meaningful representation of the lazy object for debugging
373 # without evaluating the wrapped object.
374 def __repr__(self) -> str:
375 if self._wrapped is empty:
376 repr_attr = self._setupfunc
377 else:
378 repr_attr = self._wrapped
379 return f"<{type(self).__name__}: {repr_attr!r}>"
380
381 def __copy__(self) -> SimpleLazyObject | Any:
382 if self._wrapped is empty:
383 # If uninitialized, copy the wrapper. Use SimpleLazyObject, not
384 # self.__class__, because the latter is proxied.
385 return SimpleLazyObject(self._setupfunc)
386 else:
387 # If initialized, return a copy of the wrapped object.
388 return copy.copy(self._wrapped)
389
390 def __deepcopy__(self, memo: dict[int, Any]) -> SimpleLazyObject | Any:
391 if self._wrapped is empty:
392 # We have to use SimpleLazyObject, not self.__class__, because the
393 # latter is proxied.
394 result = SimpleLazyObject(self._setupfunc)
395 memo[id(self)] = result
396 return result
397 return copy.deepcopy(self._wrapped, memo)
398
399 __add__ = new_method_proxy(operator.add)
400
401 @new_method_proxy
402 def __radd__(self: Any, other: Any) -> Any:
403 return other + self
404
405
406def partition(
407 predicate: Callable[[Any], bool], values: Any
408) -> tuple[list[Any], list[Any]]:
409 """
410 Split the values into two sets, based on the return value of the function
411 (True/False). e.g.:
412
413 >>> partition(lambda x: x > 3, range(5))
414 [0, 1, 2, 3], [4]
415 """
416 results: tuple[list[Any], list[Any]] = ([], [])
417 for item in values:
418 results[predicate(item)].append(item)
419 return results