_common.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. import os
  2. import pathlib
  3. import zipfile
  4. import tempfile
  5. import functools
  6. import contextlib
  7. def from_package(package):
  8. """
  9. Return a Traversable object for the given package.
  10. """
  11. return fallback_resources(package.__spec__)
  12. def fallback_resources(spec):
  13. package_directory = pathlib.Path(spec.origin).parent
  14. try:
  15. archive_path = spec.loader.archive
  16. rel_path = package_directory.relative_to(archive_path)
  17. return zipfile.Path(archive_path, str(rel_path) + '/')
  18. except Exception:
  19. pass
  20. return package_directory
  21. @contextlib.contextmanager
  22. def _tempfile(reader, suffix=''):
  23. # Not using tempfile.NamedTemporaryFile as it leads to deeper 'try'
  24. # blocks due to the need to close the temporary file to work on Windows
  25. # properly.
  26. fd, raw_path = tempfile.mkstemp(suffix=suffix)
  27. try:
  28. os.write(fd, reader())
  29. os.close(fd)
  30. yield pathlib.Path(raw_path)
  31. finally:
  32. try:
  33. os.remove(raw_path)
  34. except FileNotFoundError:
  35. pass
  36. @functools.singledispatch
  37. @contextlib.contextmanager
  38. def as_file(path):
  39. """
  40. Given a Traversable object, return that object as a
  41. path on the local file system in a context manager.
  42. """
  43. with _tempfile(path.read_bytes, suffix=path.name) as local:
  44. yield local
  45. @as_file.register(pathlib.Path)
  46. @contextlib.contextmanager
  47. def _(path):
  48. """
  49. Degenerate behavior for pathlib.Path objects.
  50. """
  51. yield path