_iterutils.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. from __future__ import annotations
  2. from collections.abc import Hashable, Iterable
  3. from typing import TYPE_CHECKING, Any, TypeVar, Union, overload
  4. if TYPE_CHECKING:
  5. T = TypeVar("T")
  6. HashableT = TypeVar("HashableT", bound=Hashable)
  7. ClassInfo = Union[type[T], tuple[type[T], ...]]
  8. @overload
  9. def always_list(obj: Iterable[T], base_type: ClassInfo = ...) -> list[T]: ...
  10. @overload
  11. def always_list(obj: T, base_type: ClassInfo = ...) -> list[T]: ...
  12. def always_list(obj: Any, base_type: Any = (str, bytes)) -> list[T]:
  13. """Return a guaranteed list of objects from one instance OR an iterable of such items.
  14. By default, assume the returned list should have string-like elements (`str`/`bytes`).
  15. Adapted from `more_itertools.always_iterable`, but simplified for internal use. See:
  16. https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.always_iterable
  17. """
  18. return [obj] if isinstance(obj, base_type) else list(obj)
  19. def unique_list(iterable: Iterable[HashableT]) -> list[HashableT]:
  20. """Return a deduplicated list of items from the given iterable, preserving order."""
  21. # Trick for O(1) uniqueness check that maintains order
  22. return list(dict.fromkeys(iterable))
  23. def one(
  24. iterable: Iterable[T],
  25. too_short: type[Exception] | Exception | None = None,
  26. too_long: type[Exception] | Exception | None = None,
  27. ) -> T:
  28. """Return the only item in the iterable.
  29. Note:
  30. This is intended **only** as an internal helper/convenience function,
  31. and its implementation is directly adapted from `more_itertools.one`.
  32. Users needing similar functionality are strongly encouraged to use
  33. that library instead:
  34. https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.one
  35. Args:
  36. iterable: The iterable to get the only item from.
  37. too_short: Custom exception to raise if the iterable has no items.
  38. too_long: Custom exception to raise if the iterable has multiple items.
  39. Raises:
  40. ValueError or `too_short`: If the iterable has no items.
  41. ValueError or `too_long`: If the iterable has multiple items.
  42. """
  43. # For a general iterable, avoid inadvertently iterating through all values,
  44. # which may be costly or impossible (e.g. if infinite). Only check that:
  45. # ... the first item exists
  46. it = iter(iterable)
  47. try:
  48. obj = next(it)
  49. except StopIteration:
  50. raise (too_short or ValueError("Expected 1 item in iterable, got 0")) from None
  51. # ...the second item doesn't
  52. try:
  53. _ = next(it)
  54. except StopIteration:
  55. return obj
  56. raise (
  57. too_long or ValueError("Expected 1 item in iterable, got multiple")
  58. ) from None