Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 1 | import os |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 2 | import io |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 3 | |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 4 | from . import _common |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 5 | from ._common import as_file, files |
| 6 | from .abc import ResourceReader |
| 7 | from contextlib import suppress |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 8 | from importlib.abc import ResourceLoader |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 9 | from importlib.machinery import ModuleSpec |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 10 | from io import BytesIO, TextIOWrapper |
| 11 | from pathlib import Path |
| 12 | from types import ModuleType |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 13 | from typing import ContextManager, Iterable, Union |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 14 | from typing import cast |
| 15 | from typing.io import BinaryIO, TextIO |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 16 | from collections.abc import Sequence |
| 17 | from functools import singledispatch |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 18 | |
| 19 | |
| 20 | __all__ = [ |
| 21 | 'Package', |
| 22 | 'Resource', |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 23 | 'ResourceReader', |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 24 | 'as_file', |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 25 | 'contents', |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 26 | 'files', |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 27 | 'is_resource', |
| 28 | 'open_binary', |
| 29 | 'open_text', |
| 30 | 'path', |
| 31 | 'read_binary', |
| 32 | 'read_text', |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 33 | ] |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 34 | |
| 35 | |
| 36 | Package = Union[str, ModuleType] |
| 37 | Resource = Union[str, os.PathLike] |
| 38 | |
| 39 | |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 40 | def open_binary(package: Package, resource: Resource) -> BinaryIO: |
| 41 | """Return a file-like object opened for binary reading of the resource.""" |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 42 | resource = _common.normalize_path(resource) |
| 43 | package = _common.get_package(package) |
| 44 | reader = _common.get_resource_reader(package) |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 45 | if reader is not None: |
| 46 | return reader.open_resource(resource) |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 47 | spec = cast(ModuleSpec, package.__spec__) |
| 48 | # Using pathlib doesn't work well here due to the lack of 'strict' |
| 49 | # argument for pathlib.Path.resolve() prior to Python 3.6. |
| 50 | if spec.submodule_search_locations is not None: |
| 51 | paths = spec.submodule_search_locations |
| 52 | elif spec.origin is not None: |
| 53 | paths = [os.path.dirname(os.path.abspath(spec.origin))] |
| 54 | |
| 55 | for package_path in paths: |
| 56 | full_path = os.path.join(package_path, resource) |
| 57 | try: |
| 58 | return open(full_path, mode='rb') |
| 59 | except OSError: |
| 60 | # Just assume the loader is a resource loader; all the relevant |
| 61 | # importlib.machinery loaders are and an AttributeError for |
| 62 | # get_data() will make it clear what is needed from the loader. |
| 63 | loader = cast(ResourceLoader, spec.loader) |
| 64 | data = None |
| 65 | if hasattr(spec.loader, 'get_data'): |
| 66 | with suppress(OSError): |
| 67 | data = loader.get_data(full_path) |
| 68 | if data is not None: |
| 69 | return BytesIO(data) |
| 70 | |
| 71 | raise FileNotFoundError(f'{resource!r} resource not found in {spec.name!r}') |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 72 | |
| 73 | |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 74 | def open_text( |
| 75 | package: Package, |
| 76 | resource: Resource, |
| 77 | encoding: str = 'utf-8', |
| 78 | errors: str = 'strict', |
| 79 | ) -> TextIO: |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 80 | """Return a file-like object opened for text reading of the resource.""" |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 81 | return TextIOWrapper( |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 82 | open_binary(package, resource), encoding=encoding, errors=errors |
| 83 | ) |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 84 | |
| 85 | |
| 86 | def read_binary(package: Package, resource: Resource) -> bytes: |
| 87 | """Return the binary contents of the resource.""" |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 88 | with open_binary(package, resource) as fp: |
| 89 | return fp.read() |
| 90 | |
| 91 | |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 92 | def read_text( |
| 93 | package: Package, |
| 94 | resource: Resource, |
| 95 | encoding: str = 'utf-8', |
| 96 | errors: str = 'strict', |
| 97 | ) -> str: |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 98 | """Return the decoded string of the resource. |
| 99 | |
| 100 | The decoding-related arguments have the same semantics as those of |
| 101 | bytes.decode(). |
| 102 | """ |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 103 | with open_text(package, resource, encoding, errors) as fp: |
| 104 | return fp.read() |
| 105 | |
| 106 | |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 107 | def path( |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 108 | package: Package, |
| 109 | resource: Resource, |
| 110 | ) -> 'ContextManager[Path]': |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 111 | """A context manager providing a file path object to the resource. |
| 112 | |
| 113 | If the resource does not already exist on its own on the file system, |
| 114 | a temporary file will be created. If the file was created, the file |
| 115 | will be deleted upon exiting the context manager (no exception is |
| 116 | raised if the file was deleted prior to the context manager |
| 117 | exiting). |
| 118 | """ |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 119 | reader = _common.get_resource_reader(_common.get_package(package)) |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 120 | return ( |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 121 | _path_from_reader(reader, _common.normalize_path(resource)) |
| 122 | if reader |
| 123 | else _common.as_file( |
| 124 | _common.files(package).joinpath(_common.normalize_path(resource)) |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 125 | ) |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 126 | ) |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 127 | |
| 128 | |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 129 | def _path_from_reader(reader, resource): |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 130 | return _path_from_resource_path(reader, resource) or _path_from_open_resource( |
| 131 | reader, resource |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | def _path_from_resource_path(reader, resource): |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 136 | with suppress(FileNotFoundError): |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 137 | return Path(reader.resource_path(resource)) |
| 138 | |
| 139 | |
| 140 | def _path_from_open_resource(reader, resource): |
| 141 | saved = io.BytesIO(reader.open_resource(resource).read()) |
| 142 | return _common._tempfile(saved.read, suffix=resource) |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 143 | |
| 144 | |
| 145 | def is_resource(package: Package, name: str) -> bool: |
| 146 | """True if 'name' is a resource inside 'package'. |
| 147 | |
| 148 | Directories are *not* resources. |
| 149 | """ |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 150 | package = _common.get_package(package) |
| 151 | _common.normalize_path(name) |
| 152 | reader = _common.get_resource_reader(package) |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 153 | if reader is not None: |
| 154 | return reader.is_resource(name) |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 155 | package_contents = set(contents(package)) |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 156 | if name not in package_contents: |
| 157 | return False |
Haibo Huang | 5eba2b4 | 2021-01-22 11:22:02 -0800 | [diff] [blame] | 158 | return (_common.from_package(package) / name).is_file() |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 159 | |
| 160 | |
| 161 | def contents(package: Package) -> Iterable[str]: |
| 162 | """Return an iterable of entries in 'package'. |
| 163 | |
| 164 | Note that not all entries are resources. Specifically, directories are |
| 165 | not considered resources. Use `is_resource()` on each entry returned here |
| 166 | to check if it is a resource or not. |
| 167 | """ |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 168 | package = _common.get_package(package) |
| 169 | reader = _common.get_resource_reader(package) |
Haibo Huang | d883030 | 2020-03-03 10:09:46 -0800 | [diff] [blame] | 170 | if reader is not None: |
Yi Kong | 7119932 | 2022-08-30 15:53:45 +0800 | [diff] [blame] | 171 | return _ensure_sequence(reader.contents()) |
| 172 | transversable = _common.from_package(package) |
| 173 | if transversable.is_dir(): |
| 174 | return list(item.name for item in transversable.iterdir()) |
| 175 | return [] |
| 176 | |
| 177 | |
| 178 | @singledispatch |
| 179 | def _ensure_sequence(iterable): |
| 180 | return list(iterable) |
| 181 | |
| 182 | |
| 183 | @_ensure_sequence.register(Sequence) |
| 184 | def _(iterable): |
| 185 | return iterable |