common.py 1021 B

12345678910111213141516171819202122232425262728293031323334353637383940
  1. """ common utilities """
  2. from __future__ import annotations
  3. from typing import (
  4. Any,
  5. Literal,
  6. )
  7. def _mklbl(prefix: str, n: int):
  8. return [f"{prefix}{i}" for i in range(n)]
  9. def check_indexing_smoketest_or_raises(
  10. obj,
  11. method: Literal["iloc", "loc"],
  12. key: Any,
  13. axes: Literal[0, 1] | None = None,
  14. fails=None,
  15. ) -> None:
  16. if axes is None:
  17. axes_list = [0, 1]
  18. else:
  19. assert axes in [0, 1]
  20. axes_list = [axes]
  21. for ax in axes_list:
  22. if ax < obj.ndim:
  23. # create a tuple accessor
  24. new_axes = [slice(None)] * obj.ndim
  25. new_axes[ax] = key
  26. axified = tuple(new_axes)
  27. try:
  28. getattr(obj, method).__getitem__(axified)
  29. except (IndexError, TypeError, KeyError) as detail:
  30. # if we are in fails, the ok, otherwise raise it
  31. if fails is not None:
  32. if isinstance(detail, fails):
  33. return
  34. raise