_util.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. # -*- coding: utf-8 -*-
  2. # Copyright 2018 Joshua Bronson. All Rights Reserved.
  3. #
  4. # This Source Code Form is subject to the terms of the Mozilla Public
  5. # License, v. 2.0. If a copy of the MPL was not distributed with this
  6. # file, You can obtain one at http://mozilla.org/MPL/2.0/.
  7. """Useful functions for working with bidirectional mappings and related data."""
  8. from itertools import chain, repeat
  9. from .compat import iteritems, Mapping
  10. _NULL_IT = repeat(None, 0) # repeat 0 times -> raise StopIteration from the start
  11. def _iteritems_mapping_or_iterable(arg):
  12. """Yield the items in *arg*.
  13. If *arg* is a :class:`~collections.abc.Mapping`, return an iterator over its items.
  14. Otherwise return an iterator over *arg* itself.
  15. """
  16. return iteritems(arg) if isinstance(arg, Mapping) else iter(arg)
  17. def _iteritems_args_kw(*args, **kw):
  18. """Yield the items from the positional argument (if given) and then any from *kw*.
  19. :raises TypeError: if more than one positional argument is given.
  20. """
  21. args_len = len(args)
  22. if args_len > 1:
  23. raise TypeError('Expected at most 1 positional argument, got %d' % args_len)
  24. itemchain = None
  25. if args:
  26. arg = args[0]
  27. if arg:
  28. itemchain = _iteritems_mapping_or_iterable(arg)
  29. if kw:
  30. iterkw = iteritems(kw)
  31. itemchain = chain(itemchain, iterkw) if itemchain else iterkw
  32. return itemchain or _NULL_IT
  33. def inverted(arg):
  34. """Yield the inverse items of the provided object.
  35. If *arg* has a :func:`callable` ``__inverted__`` attribute,
  36. return the result of calling it.
  37. Otherwise, return an iterator over the items in `arg`,
  38. inverting each item on the fly.
  39. *See also* :attr:`bidict.BidirectionalMapping.__inverted__`
  40. """
  41. inv = getattr(arg, '__inverted__', None)
  42. if callable(inv):
  43. return inv()
  44. return ((val, key) for (key, val) in _iteritems_mapping_or_iterable(arg))