blob: db0e0c0eeff80b5c52cd2f49b5c6e12d5c8a52c9 [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
71 raise FileNotFoundError(
72 '{!r} resource not found in {!r}'.format(resource, spec.name)
73 )
Barry Warsawdeae6b42017-12-30 15:18:06 -050074
75
Jason R. Coombs67148252021-03-04 13:43:00 -050076def open_text(
77 package: Package,
78 resource: Resource,
79 encoding: str = 'utf-8',
80 errors: str = 'strict',
81) -> TextIO:
Barry Warsawdeae6b42017-12-30 15:18:06 -050082 """Return a file-like object opened for text reading of the resource."""
Jason R. Coombs7f7e7062020-05-08 19:20:26 -040083 return TextIOWrapper(
Jason R. Coombs67148252021-03-04 13:43:00 -050084 open_binary(package, resource), encoding=encoding, errors=errors
85 )
Barry Warsawdeae6b42017-12-30 15:18:06 -050086
87
88def read_binary(package: Package, resource: Resource) -> bytes:
89 """Return the binary contents of the resource."""
Barry Warsawdeae6b42017-12-30 15:18:06 -050090 with open_binary(package, resource) as fp:
91 return fp.read()
92
93
Jason R. Coombs67148252021-03-04 13:43:00 -050094def read_text(
95 package: Package,
96 resource: Resource,
97 encoding: str = 'utf-8',
98 errors: str = 'strict',
99) -> str:
Barry Warsawdeae6b42017-12-30 15:18:06 -0500100 """Return the decoded string of the resource.
101
102 The decoding-related arguments have the same semantics as those of
103 bytes.decode().
104 """
Barry Warsawdeae6b42017-12-30 15:18:06 -0500105 with open_text(package, resource, encoding, errors) as fp:
106 return fp.read()
107
108
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400109def path(
Jason R. Coombs67148252021-03-04 13:43:00 -0500110 package: Package,
111 resource: Resource,
112) -> 'ContextManager[Path]':
Barry Warsawdeae6b42017-12-30 15:18:06 -0500113 """A context manager providing a file path object to the resource.
114
115 If the resource does not already exist on its own on the file system,
116 a temporary file will be created. If the file was created, the file
117 will be deleted upon exiting the context manager (no exception is
118 raised if the file was deleted prior to the context manager
119 exiting).
120 """
Jason R. Coombs843c2772020-06-07 21:00:51 -0400121 reader = _common.get_resource_reader(_common.get_package(package))
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400122 return (
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400123 _path_from_reader(reader, _common.normalize_path(resource))
Jason R. Coombs67148252021-03-04 13:43:00 -0500124 if reader
125 else _common.as_file(
126 _common.files(package).joinpath(_common.normalize_path(resource))
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400127 )
Jason R. Coombs67148252021-03-04 13:43:00 -0500128 )
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400129
130
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400131def _path_from_reader(reader, resource):
Jason R. Coombs67148252021-03-04 13:43:00 -0500132 return _path_from_resource_path(reader, resource) or _path_from_open_resource(
133 reader, resource
134 )
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400135
136
137def _path_from_resource_path(reader, resource):
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400138 with suppress(FileNotFoundError):
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400139 return Path(reader.resource_path(resource))
140
141
142def _path_from_open_resource(reader, resource):
143 saved = io.BytesIO(reader.open_resource(resource).read())
144 return _common._tempfile(saved.read, suffix=resource)
Barry Warsawdeae6b42017-12-30 15:18:06 -0500145
146
147def is_resource(package: Package, name: str) -> bool:
148 """True if 'name' is a resource inside 'package'.
149
150 Directories are *not* resources.
151 """
Jason R. Coombs843c2772020-06-07 21:00:51 -0400152 package = _common.get_package(package)
153 _common.normalize_path(name)
154 reader = _common.get_resource_reader(package)
Barry Warsawdeae6b42017-12-30 15:18:06 -0500155 if reader is not None:
156 return reader.is_resource(name)
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400157 package_contents = set(contents(package))
Barry Warsawdeae6b42017-12-30 15:18:06 -0500158 if name not in package_contents:
159 return False
Jason R. Coombs7f7e7062020-05-08 19:20:26 -0400160 return (_common.from_package(package) / name).is_file()
Barry Warsawdeae6b42017-12-30 15:18:06 -0500161
162
Brett Cannon3ab93652018-04-30 11:31:45 -0700163def contents(package: Package) -> Iterable[str]:
164 """Return an iterable of entries in 'package'.
Barry Warsawdeae6b42017-12-30 15:18:06 -0500165
166 Note that not all entries are resources. Specifically, directories are
167 not considered resources. Use `is_resource()` on each entry returned here
168 to check if it is a resource or not.
169 """
Jason R. Coombs843c2772020-06-07 21:00:51 -0400170 package = _common.get_package(package)
171 reader = _common.get_resource_reader(package)
Barry Warsawdeae6b42017-12-30 15:18:06 -0500172 if reader is not None:
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400173 return _ensure_sequence(reader.contents())
Jason R. Coombs67148252021-03-04 13:43:00 -0500174 transversable = _common.from_package(package)
175 if transversable.is_dir():
176 return list(item.name for item in transversable.iterdir())
177 return []
Jason R. Coombsdf8d4c82020-10-25 14:21:46 -0400178
179
180@singledispatch
181def _ensure_sequence(iterable):
182 return list(iterable)
183
184
185@_ensure_sequence.register(Sequence)
186def _(iterable):
187 return iterable