blob: 8a98663ff8e6d5e9c53b4cb292778fe2ff95fd30 [file] [log] [blame]
Barry Warsawdeae6b42017-12-30 15:18:06 -05001import os
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -04002import io
Barry Warsawdeae6b42017-12-30 15:18:06 -05003
Jason R. Coombs7f7e7062020-05-08 19:20:26 -04004from . import _common
Jason R. Coombs843c2772020-06-07 21:00:51 -04005from ._common import as_file, files
Jason R. Coombs67148252021-03-04 13:43:00 -05006from .abc import ResourceReader
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -04007from contextlib import suppress
Barry Warsawdeae6b42017-12-30 15:18:06 -05008from importlib.abc import ResourceLoader
Jason R. Coombs67148252021-03-04 13:43:00 -05009from importlib.machinery import ModuleSpec
Barry Warsawdeae6b42017-12-30 15:18:06 -050010from io import BytesIO, TextIOWrapper
11from pathlib import Path
12from types import ModuleType
Jason R. Coombs843c2772020-06-07 21:00:51 -040013from typing import ContextManager, Iterable, Union
Barry Warsawdeae6b42017-12-30 15:18:06 -050014from typing import cast
15from typing.io import BinaryIO, TextIO
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -040016from collections.abc import Sequence
17from functools import singledispatch
Barry Warsawdeae6b42017-12-30 15:18:06 -050018
19
Barry Warsaw0ed66df2018-05-17 11:41:53 -040020__all__ = [
21 'Package',
22 'Resource',
Jason R. Coombs67148252021-03-04 13:43:00 -050023 'ResourceReader',
Jason R. Coombs7f7e7062020-05-08 19:20:26 -040024 'as_file',
Barry Warsaw0ed66df2018-05-17 11:41:53 -040025 'contents',
Jason R. Coombs7f7e7062020-05-08 19:20:26 -040026 'files',
Barry Warsaw0ed66df2018-05-17 11:41:53 -040027 'is_resource',
28 'open_binary',
29 'open_text',
30 'path',
31 'read_binary',
32 'read_text',
Jason R. Coombs67148252021-03-04 13:43:00 -050033]
Barry Warsaw0ed66df2018-05-17 11:41:53 -040034
35
Barry Warsawdeae6b42017-12-30 15:18:06 -050036Package = Union[str, ModuleType]
37Resource = Union[str, os.PathLike]
38
39
Barry Warsawdeae6b42017-12-30 15:18:06 -050040def open_binary(package: Package, resource: Resource) -> BinaryIO:
41 """Return a file-like object opened for binary reading of the resource."""
Jason R. Coombs843c2772020-06-07 21:00:51 -040042 resource = _common.normalize_path(resource)
43 package = _common.get_package(package)
44 reader = _common.get_resource_reader(package)
Barry Warsawdeae6b42017-12-30 15:18:06 -050045 if reader is not None:
46 return reader.open_resource(resource)
Jason R. Coombs67148252021-03-04 13:43:00 -050047 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
Miss Islington (bot)97b45762021-05-26 14:09:27 -070071 raise FileNotFoundError(f'{resource!r} resource not found in {spec.name!r}')
Barry Warsawdeae6b42017-12-30 15:18:06 -050072
73
Jason R. Coombs67148252021-03-04 13:43:00 -050074def open_text(
75 package: Package,
76 resource: Resource,
77 encoding: str = 'utf-8',
78 errors: str = 'strict',
79) -> TextIO:
Barry Warsawdeae6b42017-12-30 15:18:06 -050080 """Return a file-like object opened for text reading of the resource."""
Jason R. Coombs7f7e7062020-05-08 19:20:26 -040081 return TextIOWrapper(
Jason R. Coombs67148252021-03-04 13:43:00 -050082 open_binary(package, resource), encoding=encoding, errors=errors
83 )
Barry Warsawdeae6b42017-12-30 15:18:06 -050084
85
86def read_binary(package: Package, resource: Resource) -> bytes:
87 """Return the binary contents of the resource."""
Barry Warsawdeae6b42017-12-30 15:18:06 -050088 with open_binary(package, resource) as fp:
89 return fp.read()
90
91
Jason R. Coombs67148252021-03-04 13:43:00 -050092def read_text(
93 package: Package,
94 resource: Resource,
95 encoding: str = 'utf-8',
96 errors: str = 'strict',
97) -> str:
Barry Warsawdeae6b42017-12-30 15:18:06 -050098 """Return the decoded string of the resource.
99
100 The decoding-related arguments have the same semantics as those of
101 bytes.decode().
102 """
Barry Warsawdeae6b42017-12-30 15:18:06 -0500103 with open_text(package, resource, encoding, errors) as fp:
104 return fp.read()
105
106
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400107def path(
Jason R. Coombs67148252021-03-04 13:43:00 -0500108 package: Package,
109 resource: Resource,
110) -> 'ContextManager[Path]':
Barry Warsawdeae6b42017-12-30 15:18:06 -0500111 """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 """
Jason R. Coombs843c2772020-06-07 21:00:51 -0400119 reader = _common.get_resource_reader(_common.get_package(package))
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400120 return (
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400121 _path_from_reader(reader, _common.normalize_path(resource))
Jason R. Coombs67148252021-03-04 13:43:00 -0500122 if reader
123 else _common.as_file(
124 _common.files(package).joinpath(_common.normalize_path(resource))
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400125 )
Jason R. Coombs67148252021-03-04 13:43:00 -0500126 )
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400127
128
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400129def _path_from_reader(reader, resource):
Jason R. Coombs67148252021-03-04 13:43:00 -0500130 return _path_from_resource_path(reader, resource) or _path_from_open_resource(
131 reader, resource
132 )
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400133
134
135def _path_from_resource_path(reader, resource):
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400136 with suppress(FileNotFoundError):
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400137 return Path(reader.resource_path(resource))
138
139
140def _path_from_open_resource(reader, resource):
141 saved = io.BytesIO(reader.open_resource(resource).read())
142 return _common._tempfile(saved.read, suffix=resource)
Barry Warsawdeae6b42017-12-30 15:18:06 -0500143
144
145def is_resource(package: Package, name: str) -> bool:
146 """True if 'name' is a resource inside 'package'.
147
148 Directories are *not* resources.
149 """
Jason R. Coombs843c2772020-06-07 21:00:51 -0400150 package = _common.get_package(package)
151 _common.normalize_path(name)
152 reader = _common.get_resource_reader(package)
Barry Warsawdeae6b42017-12-30 15:18:06 -0500153 if reader is not None:
154 return reader.is_resource(name)
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400155 package_contents = set(contents(package))
Barry Warsawdeae6b42017-12-30 15:18:06 -0500156 if name not in package_contents:
157 return False
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400158 return (_common.from_package(package) / name).is_file()
Barry Warsawdeae6b42017-12-30 15:18:06 -0500159
160
Brett Cannon3ab93652018-04-30 11:31:45 -0700161def contents(package: Package) -> Iterable[str]:
162 """Return an iterable of entries in 'package'.
Barry Warsawdeae6b42017-12-30 15:18:06 -0500163
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 """
Jason R. Coombs843c2772020-06-07 21:00:51 -0400168 package = _common.get_package(package)
169 reader = _common.get_resource_reader(package)
Barry Warsawdeae6b42017-12-30 15:18:06 -0500170 if reader is not None:
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400171 return _ensure_sequence(reader.contents())
Jason R. Coombs67148252021-03-04 13:43:00 -0500172 transversable = _common.from_package(package)
173 if transversable.is_dir():
174 return list(item.name for item in transversable.iterdir())
175 return []
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400176
177
178@singledispatch
179def _ensure_sequence(iterable):
180 return list(iterable)
181
182
183@_ensure_sequence.register(Sequence)
184def _(iterable):
185 return iterable