prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (242) hide show
  1. prefect/__init__.py +8 -56
  2. prefect/_internal/compatibility/deprecated.py +6 -115
  3. prefect/_internal/compatibility/experimental.py +4 -79
  4. prefect/_internal/concurrency/api.py +0 -34
  5. prefect/_internal/concurrency/calls.py +0 -6
  6. prefect/_internal/concurrency/cancellation.py +0 -3
  7. prefect/_internal/concurrency/event_loop.py +0 -20
  8. prefect/_internal/concurrency/inspection.py +3 -3
  9. prefect/_internal/concurrency/threads.py +35 -0
  10. prefect/_internal/concurrency/waiters.py +0 -28
  11. prefect/_internal/pydantic/__init__.py +0 -45
  12. prefect/_internal/pydantic/v1_schema.py +21 -22
  13. prefect/_internal/pydantic/v2_schema.py +0 -2
  14. prefect/_internal/pydantic/v2_validated_func.py +18 -23
  15. prefect/_internal/schemas/bases.py +44 -177
  16. prefect/_internal/schemas/fields.py +1 -43
  17. prefect/_internal/schemas/validators.py +60 -158
  18. prefect/artifacts.py +161 -14
  19. prefect/automations.py +39 -4
  20. prefect/blocks/abstract.py +1 -1
  21. prefect/blocks/core.py +268 -148
  22. prefect/blocks/fields.py +2 -57
  23. prefect/blocks/kubernetes.py +8 -12
  24. prefect/blocks/notifications.py +40 -20
  25. prefect/blocks/redis.py +168 -0
  26. prefect/blocks/system.py +22 -11
  27. prefect/blocks/webhook.py +2 -9
  28. prefect/client/base.py +4 -4
  29. prefect/client/cloud.py +8 -13
  30. prefect/client/orchestration.py +362 -340
  31. prefect/client/schemas/actions.py +92 -86
  32. prefect/client/schemas/filters.py +20 -40
  33. prefect/client/schemas/objects.py +158 -152
  34. prefect/client/schemas/responses.py +16 -24
  35. prefect/client/schemas/schedules.py +47 -35
  36. prefect/client/subscriptions.py +2 -2
  37. prefect/client/utilities.py +5 -2
  38. prefect/concurrency/asyncio.py +4 -2
  39. prefect/concurrency/events.py +1 -1
  40. prefect/concurrency/services.py +7 -4
  41. prefect/context.py +195 -27
  42. prefect/deployments/__init__.py +5 -6
  43. prefect/deployments/base.py +7 -5
  44. prefect/deployments/flow_runs.py +185 -0
  45. prefect/deployments/runner.py +50 -45
  46. prefect/deployments/schedules.py +28 -23
  47. prefect/deployments/steps/__init__.py +0 -1
  48. prefect/deployments/steps/core.py +1 -0
  49. prefect/deployments/steps/pull.py +7 -21
  50. prefect/engine.py +12 -2422
  51. prefect/events/actions.py +17 -23
  52. prefect/events/cli/automations.py +19 -6
  53. prefect/events/clients.py +14 -37
  54. prefect/events/filters.py +14 -18
  55. prefect/events/related.py +2 -2
  56. prefect/events/schemas/__init__.py +0 -5
  57. prefect/events/schemas/automations.py +55 -46
  58. prefect/events/schemas/deployment_triggers.py +7 -197
  59. prefect/events/schemas/events.py +36 -65
  60. prefect/events/schemas/labelling.py +10 -14
  61. prefect/events/utilities.py +2 -3
  62. prefect/events/worker.py +2 -3
  63. prefect/filesystems.py +6 -517
  64. prefect/{new_flow_engine.py → flow_engine.py} +315 -74
  65. prefect/flow_runs.py +379 -7
  66. prefect/flows.py +248 -165
  67. prefect/futures.py +187 -345
  68. prefect/infrastructure/__init__.py +0 -27
  69. prefect/infrastructure/provisioners/__init__.py +5 -3
  70. prefect/infrastructure/provisioners/cloud_run.py +11 -6
  71. prefect/infrastructure/provisioners/container_instance.py +11 -7
  72. prefect/infrastructure/provisioners/ecs.py +6 -4
  73. prefect/infrastructure/provisioners/modal.py +8 -5
  74. prefect/input/actions.py +2 -4
  75. prefect/input/run_input.py +9 -9
  76. prefect/logging/formatters.py +0 -2
  77. prefect/logging/handlers.py +3 -11
  78. prefect/logging/loggers.py +2 -2
  79. prefect/manifests.py +2 -1
  80. prefect/records/__init__.py +1 -0
  81. prefect/records/cache_policies.py +179 -0
  82. prefect/records/result_store.py +42 -0
  83. prefect/records/store.py +9 -0
  84. prefect/results.py +43 -39
  85. prefect/runner/runner.py +9 -9
  86. prefect/runner/server.py +6 -10
  87. prefect/runner/storage.py +3 -8
  88. prefect/runner/submit.py +2 -2
  89. prefect/runner/utils.py +2 -2
  90. prefect/serializers.py +24 -35
  91. prefect/server/api/collections_data/views/aggregate-worker-metadata.json +5 -14
  92. prefect/settings.py +76 -136
  93. prefect/states.py +22 -50
  94. prefect/task_engine.py +666 -56
  95. prefect/task_runners.py +272 -300
  96. prefect/task_runs.py +203 -0
  97. prefect/{task_server.py → task_worker.py} +89 -60
  98. prefect/tasks.py +358 -341
  99. prefect/transactions.py +224 -0
  100. prefect/types/__init__.py +61 -82
  101. prefect/utilities/asyncutils.py +195 -136
  102. prefect/utilities/callables.py +121 -41
  103. prefect/utilities/collections.py +23 -38
  104. prefect/utilities/dispatch.py +11 -3
  105. prefect/utilities/dockerutils.py +4 -0
  106. prefect/utilities/engine.py +140 -20
  107. prefect/utilities/importtools.py +26 -27
  108. prefect/utilities/pydantic.py +128 -38
  109. prefect/utilities/schema_tools/hydration.py +5 -1
  110. prefect/utilities/templating.py +12 -2
  111. prefect/variables.py +84 -62
  112. prefect/workers/__init__.py +0 -1
  113. prefect/workers/base.py +26 -18
  114. prefect/workers/process.py +3 -8
  115. prefect/workers/server.py +2 -2
  116. {prefect_client-2.19.4.dist-info → prefect_client-3.0.0rc2.dist-info}/METADATA +23 -21
  117. prefect_client-3.0.0rc2.dist-info/RECORD +179 -0
  118. prefect/_internal/pydantic/_base_model.py +0 -51
  119. prefect/_internal/pydantic/_compat.py +0 -82
  120. prefect/_internal/pydantic/_flags.py +0 -20
  121. prefect/_internal/pydantic/_types.py +0 -8
  122. prefect/_internal/pydantic/utilities/__init__.py +0 -0
  123. prefect/_internal/pydantic/utilities/config_dict.py +0 -72
  124. prefect/_internal/pydantic/utilities/field_validator.py +0 -150
  125. prefect/_internal/pydantic/utilities/model_construct.py +0 -56
  126. prefect/_internal/pydantic/utilities/model_copy.py +0 -55
  127. prefect/_internal/pydantic/utilities/model_dump.py +0 -136
  128. prefect/_internal/pydantic/utilities/model_dump_json.py +0 -112
  129. prefect/_internal/pydantic/utilities/model_fields.py +0 -50
  130. prefect/_internal/pydantic/utilities/model_fields_set.py +0 -29
  131. prefect/_internal/pydantic/utilities/model_json_schema.py +0 -82
  132. prefect/_internal/pydantic/utilities/model_rebuild.py +0 -80
  133. prefect/_internal/pydantic/utilities/model_validate.py +0 -75
  134. prefect/_internal/pydantic/utilities/model_validate_json.py +0 -68
  135. prefect/_internal/pydantic/utilities/model_validator.py +0 -87
  136. prefect/_internal/pydantic/utilities/type_adapter.py +0 -71
  137. prefect/_vendor/__init__.py +0 -0
  138. prefect/_vendor/fastapi/__init__.py +0 -25
  139. prefect/_vendor/fastapi/applications.py +0 -946
  140. prefect/_vendor/fastapi/background.py +0 -3
  141. prefect/_vendor/fastapi/concurrency.py +0 -44
  142. prefect/_vendor/fastapi/datastructures.py +0 -58
  143. prefect/_vendor/fastapi/dependencies/__init__.py +0 -0
  144. prefect/_vendor/fastapi/dependencies/models.py +0 -64
  145. prefect/_vendor/fastapi/dependencies/utils.py +0 -877
  146. prefect/_vendor/fastapi/encoders.py +0 -177
  147. prefect/_vendor/fastapi/exception_handlers.py +0 -40
  148. prefect/_vendor/fastapi/exceptions.py +0 -46
  149. prefect/_vendor/fastapi/logger.py +0 -3
  150. prefect/_vendor/fastapi/middleware/__init__.py +0 -1
  151. prefect/_vendor/fastapi/middleware/asyncexitstack.py +0 -25
  152. prefect/_vendor/fastapi/middleware/cors.py +0 -3
  153. prefect/_vendor/fastapi/middleware/gzip.py +0 -3
  154. prefect/_vendor/fastapi/middleware/httpsredirect.py +0 -3
  155. prefect/_vendor/fastapi/middleware/trustedhost.py +0 -3
  156. prefect/_vendor/fastapi/middleware/wsgi.py +0 -3
  157. prefect/_vendor/fastapi/openapi/__init__.py +0 -0
  158. prefect/_vendor/fastapi/openapi/constants.py +0 -2
  159. prefect/_vendor/fastapi/openapi/docs.py +0 -203
  160. prefect/_vendor/fastapi/openapi/models.py +0 -480
  161. prefect/_vendor/fastapi/openapi/utils.py +0 -485
  162. prefect/_vendor/fastapi/param_functions.py +0 -340
  163. prefect/_vendor/fastapi/params.py +0 -453
  164. prefect/_vendor/fastapi/requests.py +0 -4
  165. prefect/_vendor/fastapi/responses.py +0 -40
  166. prefect/_vendor/fastapi/routing.py +0 -1331
  167. prefect/_vendor/fastapi/security/__init__.py +0 -15
  168. prefect/_vendor/fastapi/security/api_key.py +0 -98
  169. prefect/_vendor/fastapi/security/base.py +0 -6
  170. prefect/_vendor/fastapi/security/http.py +0 -172
  171. prefect/_vendor/fastapi/security/oauth2.py +0 -227
  172. prefect/_vendor/fastapi/security/open_id_connect_url.py +0 -34
  173. prefect/_vendor/fastapi/security/utils.py +0 -10
  174. prefect/_vendor/fastapi/staticfiles.py +0 -1
  175. prefect/_vendor/fastapi/templating.py +0 -3
  176. prefect/_vendor/fastapi/testclient.py +0 -1
  177. prefect/_vendor/fastapi/types.py +0 -3
  178. prefect/_vendor/fastapi/utils.py +0 -235
  179. prefect/_vendor/fastapi/websockets.py +0 -7
  180. prefect/_vendor/starlette/__init__.py +0 -1
  181. prefect/_vendor/starlette/_compat.py +0 -28
  182. prefect/_vendor/starlette/_exception_handler.py +0 -80
  183. prefect/_vendor/starlette/_utils.py +0 -88
  184. prefect/_vendor/starlette/applications.py +0 -261
  185. prefect/_vendor/starlette/authentication.py +0 -159
  186. prefect/_vendor/starlette/background.py +0 -43
  187. prefect/_vendor/starlette/concurrency.py +0 -59
  188. prefect/_vendor/starlette/config.py +0 -151
  189. prefect/_vendor/starlette/convertors.py +0 -87
  190. prefect/_vendor/starlette/datastructures.py +0 -707
  191. prefect/_vendor/starlette/endpoints.py +0 -130
  192. prefect/_vendor/starlette/exceptions.py +0 -60
  193. prefect/_vendor/starlette/formparsers.py +0 -276
  194. prefect/_vendor/starlette/middleware/__init__.py +0 -17
  195. prefect/_vendor/starlette/middleware/authentication.py +0 -52
  196. prefect/_vendor/starlette/middleware/base.py +0 -220
  197. prefect/_vendor/starlette/middleware/cors.py +0 -176
  198. prefect/_vendor/starlette/middleware/errors.py +0 -265
  199. prefect/_vendor/starlette/middleware/exceptions.py +0 -74
  200. prefect/_vendor/starlette/middleware/gzip.py +0 -113
  201. prefect/_vendor/starlette/middleware/httpsredirect.py +0 -19
  202. prefect/_vendor/starlette/middleware/sessions.py +0 -82
  203. prefect/_vendor/starlette/middleware/trustedhost.py +0 -64
  204. prefect/_vendor/starlette/middleware/wsgi.py +0 -147
  205. prefect/_vendor/starlette/requests.py +0 -328
  206. prefect/_vendor/starlette/responses.py +0 -347
  207. prefect/_vendor/starlette/routing.py +0 -933
  208. prefect/_vendor/starlette/schemas.py +0 -154
  209. prefect/_vendor/starlette/staticfiles.py +0 -248
  210. prefect/_vendor/starlette/status.py +0 -199
  211. prefect/_vendor/starlette/templating.py +0 -231
  212. prefect/_vendor/starlette/testclient.py +0 -804
  213. prefect/_vendor/starlette/types.py +0 -30
  214. prefect/_vendor/starlette/websockets.py +0 -193
  215. prefect/agent.py +0 -698
  216. prefect/deployments/deployments.py +0 -1042
  217. prefect/deprecated/__init__.py +0 -0
  218. prefect/deprecated/data_documents.py +0 -350
  219. prefect/deprecated/packaging/__init__.py +0 -12
  220. prefect/deprecated/packaging/base.py +0 -96
  221. prefect/deprecated/packaging/docker.py +0 -146
  222. prefect/deprecated/packaging/file.py +0 -92
  223. prefect/deprecated/packaging/orion.py +0 -80
  224. prefect/deprecated/packaging/serializers.py +0 -171
  225. prefect/events/instrument.py +0 -135
  226. prefect/infrastructure/base.py +0 -323
  227. prefect/infrastructure/container.py +0 -818
  228. prefect/infrastructure/kubernetes.py +0 -920
  229. prefect/infrastructure/process.py +0 -289
  230. prefect/new_task_engine.py +0 -423
  231. prefect/pydantic/__init__.py +0 -76
  232. prefect/pydantic/main.py +0 -39
  233. prefect/software/__init__.py +0 -2
  234. prefect/software/base.py +0 -50
  235. prefect/software/conda.py +0 -199
  236. prefect/software/pip.py +0 -122
  237. prefect/software/python.py +0 -52
  238. prefect/workers/block.py +0 -218
  239. prefect_client-2.19.4.dist-info/RECORD +0 -292
  240. {prefect_client-2.19.4.dist-info → prefect_client-3.0.0rc2.dist-info}/LICENSE +0 -0
  241. {prefect_client-2.19.4.dist-info → prefect_client-3.0.0rc2.dist-info}/WHEEL +0 -0
  242. {prefect_client-2.19.4.dist-info → prefect_client-3.0.0rc2.dist-info}/top_level.txt +0 -0
@@ -1,804 +0,0 @@
1
- import contextlib
2
- import inspect
3
- import io
4
- import json
5
- import math
6
- import queue
7
- import typing
8
- import warnings
9
- from concurrent.futures import Future
10
- from types import GeneratorType
11
- from urllib.parse import unquote, urljoin
12
-
13
- import anyio
14
- import anyio.from_thread
15
- from anyio.abc import ObjectReceiveStream, ObjectSendStream
16
- from anyio.streams.stapled import StapledObjectStream
17
- from prefect._vendor.starlette._utils import is_async_callable
18
- from prefect._vendor.starlette.types import ASGIApp, Message, Receive, Scope, Send
19
- from prefect._vendor.starlette.websockets import WebSocketDisconnect
20
-
21
- try:
22
- import httpx
23
- except ModuleNotFoundError: # pragma: no cover
24
- raise RuntimeError(
25
- "The starlette.testclient module requires the httpx package to be installed.\n"
26
- "You can install this with:\n"
27
- " $ pip install httpx\n"
28
- )
29
- _PortalFactoryType = typing.Callable[
30
- [], typing.ContextManager[anyio.abc.BlockingPortal]
31
- ]
32
-
33
- ASGIInstance = typing.Callable[[Receive, Send], typing.Awaitable[None]]
34
- ASGI2App = typing.Callable[[Scope], ASGIInstance]
35
- ASGI3App = typing.Callable[[Scope, Receive, Send], typing.Awaitable[None]]
36
-
37
-
38
- _RequestData = typing.Mapping[str, typing.Union[str, typing.Iterable[str]]]
39
-
40
-
41
- def _is_asgi3(app: typing.Union[ASGI2App, ASGI3App]) -> bool:
42
- if inspect.isclass(app):
43
- return hasattr(app, "__await__")
44
- return is_async_callable(app)
45
-
46
-
47
- class _WrapASGI2:
48
- """
49
- Provide an ASGI3 interface onto an ASGI2 app.
50
- """
51
-
52
- def __init__(self, app: ASGI2App) -> None:
53
- self.app = app
54
-
55
- async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
56
- instance = self.app(scope)
57
- await instance(receive, send)
58
-
59
-
60
- class _AsyncBackend(typing.TypedDict):
61
- backend: str
62
- backend_options: typing.Dict[str, typing.Any]
63
-
64
-
65
- class _Upgrade(Exception):
66
- def __init__(self, session: "WebSocketTestSession") -> None:
67
- self.session = session
68
-
69
-
70
- class WebSocketTestSession:
71
- def __init__(
72
- self,
73
- app: ASGI3App,
74
- scope: Scope,
75
- portal_factory: _PortalFactoryType,
76
- ) -> None:
77
- self.app = app
78
- self.scope = scope
79
- self.accepted_subprotocol = None
80
- self.portal_factory = portal_factory
81
- self._receive_queue: "queue.Queue[Message]" = queue.Queue()
82
- self._send_queue: "queue.Queue[Message | BaseException]" = queue.Queue()
83
- self.extra_headers = None
84
-
85
- def __enter__(self) -> "WebSocketTestSession":
86
- self.exit_stack = contextlib.ExitStack()
87
- self.portal = self.exit_stack.enter_context(self.portal_factory())
88
-
89
- try:
90
- _: "Future[None]" = self.portal.start_task_soon(self._run)
91
- self.send({"type": "websocket.connect"})
92
- message = self.receive()
93
- self._raise_on_close(message)
94
- except Exception:
95
- self.exit_stack.close()
96
- raise
97
- self.accepted_subprotocol = message.get("subprotocol", None)
98
- self.extra_headers = message.get("headers", None)
99
- return self
100
-
101
- def __exit__(self, *args: typing.Any) -> None:
102
- try:
103
- self.close(1000)
104
- finally:
105
- self.exit_stack.close()
106
- while not self._send_queue.empty():
107
- message = self._send_queue.get()
108
- if isinstance(message, BaseException):
109
- raise message
110
-
111
- async def _run(self) -> None:
112
- """
113
- The sub-thread in which the websocket session runs.
114
- """
115
- scope = self.scope
116
- receive = self._asgi_receive
117
- send = self._asgi_send
118
- try:
119
- await self.app(scope, receive, send)
120
- except BaseException as exc:
121
- self._send_queue.put(exc)
122
- raise
123
-
124
- async def _asgi_receive(self) -> Message:
125
- while self._receive_queue.empty():
126
- await anyio.sleep(0)
127
- return self._receive_queue.get()
128
-
129
- async def _asgi_send(self, message: Message) -> None:
130
- self._send_queue.put(message)
131
-
132
- def _raise_on_close(self, message: Message) -> None:
133
- if message["type"] == "websocket.close":
134
- raise WebSocketDisconnect(
135
- message.get("code", 1000), message.get("reason", "")
136
- )
137
-
138
- def send(self, message: Message) -> None:
139
- self._receive_queue.put(message)
140
-
141
- def send_text(self, data: str) -> None:
142
- self.send({"type": "websocket.receive", "text": data})
143
-
144
- def send_bytes(self, data: bytes) -> None:
145
- self.send({"type": "websocket.receive", "bytes": data})
146
-
147
- def send_json(self, data: typing.Any, mode: str = "text") -> None:
148
- assert mode in ["text", "binary"]
149
- text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
150
- if mode == "text":
151
- self.send({"type": "websocket.receive", "text": text})
152
- else:
153
- self.send({"type": "websocket.receive", "bytes": text.encode("utf-8")})
154
-
155
- def close(self, code: int = 1000, reason: typing.Union[str, None] = None) -> None:
156
- self.send({"type": "websocket.disconnect", "code": code, "reason": reason})
157
-
158
- def receive(self) -> Message:
159
- message = self._send_queue.get()
160
- if isinstance(message, BaseException):
161
- raise message
162
- return message
163
-
164
- def receive_text(self) -> str:
165
- message = self.receive()
166
- self._raise_on_close(message)
167
- return typing.cast(str, message["text"])
168
-
169
- def receive_bytes(self) -> bytes:
170
- message = self.receive()
171
- self._raise_on_close(message)
172
- return typing.cast(bytes, message["bytes"])
173
-
174
- def receive_json(self, mode: str = "text") -> typing.Any:
175
- assert mode in ["text", "binary"]
176
- message = self.receive()
177
- self._raise_on_close(message)
178
- if mode == "text":
179
- text = message["text"]
180
- else:
181
- text = message["bytes"].decode("utf-8")
182
- return json.loads(text)
183
-
184
-
185
- class _TestClientTransport(httpx.BaseTransport):
186
- def __init__(
187
- self,
188
- app: ASGI3App,
189
- portal_factory: _PortalFactoryType,
190
- raise_server_exceptions: bool = True,
191
- root_path: str = "",
192
- *,
193
- app_state: typing.Dict[str, typing.Any],
194
- ) -> None:
195
- self.app = app
196
- self.raise_server_exceptions = raise_server_exceptions
197
- self.root_path = root_path
198
- self.portal_factory = portal_factory
199
- self.app_state = app_state
200
-
201
- def handle_request(self, request: httpx.Request) -> httpx.Response:
202
- scheme = request.url.scheme
203
- netloc = request.url.netloc.decode(encoding="ascii")
204
- path = request.url.path
205
- raw_path = request.url.raw_path
206
- query = request.url.query.decode(encoding="ascii")
207
-
208
- default_port = {"http": 80, "ws": 80, "https": 443, "wss": 443}[scheme]
209
-
210
- if ":" in netloc:
211
- host, port_string = netloc.split(":", 1)
212
- port = int(port_string)
213
- else:
214
- host = netloc
215
- port = default_port
216
-
217
- # Include the 'host' header.
218
- if "host" in request.headers:
219
- headers: typing.List[typing.Tuple[bytes, bytes]] = []
220
- elif port == default_port: # pragma: no cover
221
- headers = [(b"host", host.encode())]
222
- else: # pragma: no cover
223
- headers = [(b"host", (f"{host}:{port}").encode())]
224
-
225
- # Include other request headers.
226
- headers += [
227
- (key.lower().encode(), value.encode())
228
- for key, value in request.headers.multi_items()
229
- ]
230
-
231
- scope: typing.Dict[str, typing.Any]
232
-
233
- if scheme in {"ws", "wss"}:
234
- subprotocol = request.headers.get("sec-websocket-protocol", None)
235
- if subprotocol is None:
236
- subprotocols: typing.Sequence[str] = []
237
- else:
238
- subprotocols = [value.strip() for value in subprotocol.split(",")]
239
- scope = {
240
- "type": "websocket",
241
- "path": unquote(path),
242
- "raw_path": raw_path,
243
- "root_path": self.root_path,
244
- "scheme": scheme,
245
- "query_string": query.encode(),
246
- "headers": headers,
247
- "client": ["testclient", 50000],
248
- "server": [host, port],
249
- "subprotocols": subprotocols,
250
- "state": self.app_state.copy(),
251
- }
252
- session = WebSocketTestSession(self.app, scope, self.portal_factory)
253
- raise _Upgrade(session)
254
-
255
- scope = {
256
- "type": "http",
257
- "http_version": "1.1",
258
- "method": request.method,
259
- "path": unquote(path),
260
- "raw_path": raw_path,
261
- "root_path": self.root_path,
262
- "scheme": scheme,
263
- "query_string": query.encode(),
264
- "headers": headers,
265
- "client": ["testclient", 50000],
266
- "server": [host, port],
267
- "extensions": {"http.response.debug": {}},
268
- "state": self.app_state.copy(),
269
- }
270
-
271
- request_complete = False
272
- response_started = False
273
- response_complete: anyio.Event
274
- raw_kwargs: typing.Dict[str, typing.Any] = {"stream": io.BytesIO()}
275
- template = None
276
- context = None
277
-
278
- async def receive() -> Message:
279
- nonlocal request_complete
280
-
281
- if request_complete:
282
- if not response_complete.is_set():
283
- await response_complete.wait()
284
- return {"type": "http.disconnect"}
285
-
286
- body = request.read()
287
- if isinstance(body, str):
288
- body_bytes: bytes = body.encode("utf-8") # pragma: no cover
289
- elif body is None:
290
- body_bytes = b"" # pragma: no cover
291
- elif isinstance(body, GeneratorType):
292
- try: # pragma: no cover
293
- chunk = body.send(None)
294
- if isinstance(chunk, str):
295
- chunk = chunk.encode("utf-8")
296
- return {"type": "http.request", "body": chunk, "more_body": True}
297
- except StopIteration: # pragma: no cover
298
- request_complete = True
299
- return {"type": "http.request", "body": b""}
300
- else:
301
- body_bytes = body
302
-
303
- request_complete = True
304
- return {"type": "http.request", "body": body_bytes}
305
-
306
- async def send(message: Message) -> None:
307
- nonlocal raw_kwargs, response_started, template, context
308
-
309
- if message["type"] == "http.response.start":
310
- assert (
311
- not response_started
312
- ), 'Received multiple "http.response.start" messages.'
313
- raw_kwargs["status_code"] = message["status"]
314
- raw_kwargs["headers"] = [
315
- (key.decode(), value.decode())
316
- for key, value in message.get("headers", [])
317
- ]
318
- response_started = True
319
- elif message["type"] == "http.response.body":
320
- assert (
321
- response_started
322
- ), 'Received "http.response.body" without "http.response.start".'
323
- assert (
324
- not response_complete.is_set()
325
- ), 'Received "http.response.body" after response completed.'
326
- body = message.get("body", b"")
327
- more_body = message.get("more_body", False)
328
- if request.method != "HEAD":
329
- raw_kwargs["stream"].write(body)
330
- if not more_body:
331
- raw_kwargs["stream"].seek(0)
332
- response_complete.set()
333
- elif message["type"] == "http.response.debug":
334
- template = message["info"]["template"]
335
- context = message["info"]["context"]
336
-
337
- try:
338
- with self.portal_factory() as portal:
339
- response_complete = portal.call(anyio.Event)
340
- portal.call(self.app, scope, receive, send)
341
- except BaseException as exc:
342
- if self.raise_server_exceptions:
343
- raise exc
344
-
345
- if self.raise_server_exceptions:
346
- assert response_started, "TestClient did not receive any response."
347
- elif not response_started:
348
- raw_kwargs = {
349
- "status_code": 500,
350
- "headers": [],
351
- "stream": io.BytesIO(),
352
- }
353
-
354
- raw_kwargs["stream"] = httpx.ByteStream(raw_kwargs["stream"].read())
355
-
356
- response = httpx.Response(**raw_kwargs, request=request)
357
- if template is not None:
358
- response.template = template # type: ignore[attr-defined]
359
- response.context = context # type: ignore[attr-defined]
360
- return response
361
-
362
-
363
- class TestClient(httpx.Client):
364
- __test__ = False
365
- task: "Future[None]"
366
- portal: typing.Optional[anyio.abc.BlockingPortal] = None
367
-
368
- def __init__(
369
- self,
370
- app: ASGIApp,
371
- base_url: str = "http://testserver",
372
- raise_server_exceptions: bool = True,
373
- root_path: str = "",
374
- backend: str = "asyncio",
375
- backend_options: typing.Optional[typing.Dict[str, typing.Any]] = None,
376
- cookies: httpx._types.CookieTypes = None,
377
- headers: typing.Dict[str, str] = None,
378
- follow_redirects: bool = True,
379
- ) -> None:
380
- self.async_backend = _AsyncBackend(
381
- backend=backend, backend_options=backend_options or {}
382
- )
383
- if _is_asgi3(app):
384
- app = typing.cast(ASGI3App, app)
385
- asgi_app = app
386
- else:
387
- app = typing.cast(ASGI2App, app) # type: ignore[assignment]
388
- asgi_app = _WrapASGI2(app) # type: ignore[arg-type]
389
- self.app = asgi_app
390
- self.app_state: typing.Dict[str, typing.Any] = {}
391
- transport = _TestClientTransport(
392
- self.app,
393
- portal_factory=self._portal_factory,
394
- raise_server_exceptions=raise_server_exceptions,
395
- root_path=root_path,
396
- app_state=self.app_state,
397
- )
398
- if headers is None:
399
- headers = {}
400
- headers.setdefault("user-agent", "testclient")
401
- super().__init__(
402
- base_url=base_url,
403
- headers=headers,
404
- transport=transport,
405
- follow_redirects=follow_redirects,
406
- cookies=cookies,
407
- )
408
-
409
- @contextlib.contextmanager
410
- def _portal_factory(self) -> typing.Generator[anyio.abc.BlockingPortal, None, None]:
411
- if self.portal is not None:
412
- yield self.portal
413
- else:
414
- with anyio.from_thread.start_blocking_portal(
415
- **self.async_backend
416
- ) as portal:
417
- yield portal
418
-
419
- def _choose_redirect_arg(
420
- self,
421
- follow_redirects: typing.Optional[bool],
422
- allow_redirects: typing.Optional[bool],
423
- ) -> typing.Union[bool, httpx._client.UseClientDefault]:
424
- redirect: typing.Union[
425
- bool, httpx._client.UseClientDefault
426
- ] = httpx._client.USE_CLIENT_DEFAULT
427
- if allow_redirects is not None:
428
- message = (
429
- "The `allow_redirects` argument is deprecated. "
430
- "Use `follow_redirects` instead."
431
- )
432
- warnings.warn(message, DeprecationWarning)
433
- redirect = allow_redirects
434
- if follow_redirects is not None:
435
- redirect = follow_redirects
436
- elif allow_redirects is not None and follow_redirects is not None:
437
- raise RuntimeError( # pragma: no cover
438
- "Cannot use both `allow_redirects` and `follow_redirects`."
439
- )
440
- return redirect
441
-
442
- def request( # type: ignore[override]
443
- self,
444
- method: str,
445
- url: httpx._types.URLTypes,
446
- *,
447
- content: typing.Optional[httpx._types.RequestContent] = None,
448
- data: typing.Optional[_RequestData] = None,
449
- files: typing.Optional[httpx._types.RequestFiles] = None,
450
- json: typing.Any = None,
451
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
452
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
453
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
454
- auth: typing.Union[
455
- httpx._types.AuthTypes, httpx._client.UseClientDefault
456
- ] = httpx._client.USE_CLIENT_DEFAULT,
457
- follow_redirects: typing.Optional[bool] = None,
458
- allow_redirects: typing.Optional[bool] = None,
459
- timeout: typing.Union[
460
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
461
- ] = httpx._client.USE_CLIENT_DEFAULT,
462
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
463
- ) -> httpx.Response:
464
- url = self._merge_url(url)
465
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
466
- return super().request(
467
- method,
468
- url,
469
- content=content,
470
- data=data,
471
- files=files,
472
- json=json,
473
- params=params,
474
- headers=headers,
475
- cookies=cookies,
476
- auth=auth,
477
- follow_redirects=redirect,
478
- timeout=timeout,
479
- extensions=extensions,
480
- )
481
-
482
- def get( # type: ignore[override]
483
- self,
484
- url: httpx._types.URLTypes,
485
- *,
486
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
487
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
488
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
489
- auth: typing.Union[
490
- httpx._types.AuthTypes, httpx._client.UseClientDefault
491
- ] = httpx._client.USE_CLIENT_DEFAULT,
492
- follow_redirects: typing.Optional[bool] = None,
493
- allow_redirects: typing.Optional[bool] = None,
494
- timeout: typing.Union[
495
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
496
- ] = httpx._client.USE_CLIENT_DEFAULT,
497
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
498
- ) -> httpx.Response:
499
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
500
- return super().get(
501
- url,
502
- params=params,
503
- headers=headers,
504
- cookies=cookies,
505
- auth=auth,
506
- follow_redirects=redirect,
507
- timeout=timeout,
508
- extensions=extensions,
509
- )
510
-
511
- def options( # type: ignore[override]
512
- self,
513
- url: httpx._types.URLTypes,
514
- *,
515
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
516
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
517
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
518
- auth: typing.Union[
519
- httpx._types.AuthTypes, httpx._client.UseClientDefault
520
- ] = httpx._client.USE_CLIENT_DEFAULT,
521
- follow_redirects: typing.Optional[bool] = None,
522
- allow_redirects: typing.Optional[bool] = None,
523
- timeout: typing.Union[
524
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
525
- ] = httpx._client.USE_CLIENT_DEFAULT,
526
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
527
- ) -> httpx.Response:
528
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
529
- return super().options(
530
- url,
531
- params=params,
532
- headers=headers,
533
- cookies=cookies,
534
- auth=auth,
535
- follow_redirects=redirect,
536
- timeout=timeout,
537
- extensions=extensions,
538
- )
539
-
540
- def head( # type: ignore[override]
541
- self,
542
- url: httpx._types.URLTypes,
543
- *,
544
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
545
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
546
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
547
- auth: typing.Union[
548
- httpx._types.AuthTypes, httpx._client.UseClientDefault
549
- ] = httpx._client.USE_CLIENT_DEFAULT,
550
- follow_redirects: typing.Optional[bool] = None,
551
- allow_redirects: typing.Optional[bool] = None,
552
- timeout: typing.Union[
553
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
554
- ] = httpx._client.USE_CLIENT_DEFAULT,
555
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
556
- ) -> httpx.Response:
557
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
558
- return super().head(
559
- url,
560
- params=params,
561
- headers=headers,
562
- cookies=cookies,
563
- auth=auth,
564
- follow_redirects=redirect,
565
- timeout=timeout,
566
- extensions=extensions,
567
- )
568
-
569
- def post( # type: ignore[override]
570
- self,
571
- url: httpx._types.URLTypes,
572
- *,
573
- content: typing.Optional[httpx._types.RequestContent] = None,
574
- data: typing.Optional[_RequestData] = None,
575
- files: typing.Optional[httpx._types.RequestFiles] = None,
576
- json: typing.Any = None,
577
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
578
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
579
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
580
- auth: typing.Union[
581
- httpx._types.AuthTypes, httpx._client.UseClientDefault
582
- ] = httpx._client.USE_CLIENT_DEFAULT,
583
- follow_redirects: typing.Optional[bool] = None,
584
- allow_redirects: typing.Optional[bool] = None,
585
- timeout: typing.Union[
586
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
587
- ] = httpx._client.USE_CLIENT_DEFAULT,
588
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
589
- ) -> httpx.Response:
590
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
591
- return super().post(
592
- url,
593
- content=content,
594
- data=data,
595
- files=files,
596
- json=json,
597
- params=params,
598
- headers=headers,
599
- cookies=cookies,
600
- auth=auth,
601
- follow_redirects=redirect,
602
- timeout=timeout,
603
- extensions=extensions,
604
- )
605
-
606
- def put( # type: ignore[override]
607
- self,
608
- url: httpx._types.URLTypes,
609
- *,
610
- content: typing.Optional[httpx._types.RequestContent] = None,
611
- data: typing.Optional[_RequestData] = None,
612
- files: typing.Optional[httpx._types.RequestFiles] = None,
613
- json: typing.Any = None,
614
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
615
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
616
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
617
- auth: typing.Union[
618
- httpx._types.AuthTypes, httpx._client.UseClientDefault
619
- ] = httpx._client.USE_CLIENT_DEFAULT,
620
- follow_redirects: typing.Optional[bool] = None,
621
- allow_redirects: typing.Optional[bool] = None,
622
- timeout: typing.Union[
623
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
624
- ] = httpx._client.USE_CLIENT_DEFAULT,
625
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
626
- ) -> httpx.Response:
627
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
628
- return super().put(
629
- url,
630
- content=content,
631
- data=data,
632
- files=files,
633
- json=json,
634
- params=params,
635
- headers=headers,
636
- cookies=cookies,
637
- auth=auth,
638
- follow_redirects=redirect,
639
- timeout=timeout,
640
- extensions=extensions,
641
- )
642
-
643
- def patch( # type: ignore[override]
644
- self,
645
- url: httpx._types.URLTypes,
646
- *,
647
- content: typing.Optional[httpx._types.RequestContent] = None,
648
- data: typing.Optional[_RequestData] = None,
649
- files: typing.Optional[httpx._types.RequestFiles] = None,
650
- json: typing.Any = None,
651
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
652
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
653
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
654
- auth: typing.Union[
655
- httpx._types.AuthTypes, httpx._client.UseClientDefault
656
- ] = httpx._client.USE_CLIENT_DEFAULT,
657
- follow_redirects: typing.Optional[bool] = None,
658
- allow_redirects: typing.Optional[bool] = None,
659
- timeout: typing.Union[
660
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
661
- ] = httpx._client.USE_CLIENT_DEFAULT,
662
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
663
- ) -> httpx.Response:
664
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
665
- return super().patch(
666
- url,
667
- content=content,
668
- data=data,
669
- files=files,
670
- json=json,
671
- params=params,
672
- headers=headers,
673
- cookies=cookies,
674
- auth=auth,
675
- follow_redirects=redirect,
676
- timeout=timeout,
677
- extensions=extensions,
678
- )
679
-
680
- def delete( # type: ignore[override]
681
- self,
682
- url: httpx._types.URLTypes,
683
- *,
684
- params: typing.Optional[httpx._types.QueryParamTypes] = None,
685
- headers: typing.Optional[httpx._types.HeaderTypes] = None,
686
- cookies: typing.Optional[httpx._types.CookieTypes] = None,
687
- auth: typing.Union[
688
- httpx._types.AuthTypes, httpx._client.UseClientDefault
689
- ] = httpx._client.USE_CLIENT_DEFAULT,
690
- follow_redirects: typing.Optional[bool] = None,
691
- allow_redirects: typing.Optional[bool] = None,
692
- timeout: typing.Union[
693
- httpx._types.TimeoutTypes, httpx._client.UseClientDefault
694
- ] = httpx._client.USE_CLIENT_DEFAULT,
695
- extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
696
- ) -> httpx.Response:
697
- redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
698
- return super().delete(
699
- url,
700
- params=params,
701
- headers=headers,
702
- cookies=cookies,
703
- auth=auth,
704
- follow_redirects=redirect,
705
- timeout=timeout,
706
- extensions=extensions,
707
- )
708
-
709
- def websocket_connect(
710
- self, url: str, subprotocols: typing.Sequence[str] = None, **kwargs: typing.Any
711
- ) -> "WebSocketTestSession":
712
- url = urljoin("ws://testserver", url)
713
- headers = kwargs.get("headers", {})
714
- headers.setdefault("connection", "upgrade")
715
- headers.setdefault("sec-websocket-key", "testserver==")
716
- headers.setdefault("sec-websocket-version", "13")
717
- if subprotocols is not None:
718
- headers.setdefault("sec-websocket-protocol", ", ".join(subprotocols))
719
- kwargs["headers"] = headers
720
- try:
721
- super().request("GET", url, **kwargs)
722
- except _Upgrade as exc:
723
- session = exc.session
724
- else:
725
- raise RuntimeError("Expected WebSocket upgrade") # pragma: no cover
726
-
727
- return session
728
-
729
- def __enter__(self) -> "TestClient":
730
- with contextlib.ExitStack() as stack:
731
- self.portal = portal = stack.enter_context(
732
- anyio.from_thread.start_blocking_portal(**self.async_backend)
733
- )
734
-
735
- @stack.callback
736
- def reset_portal() -> None:
737
- self.portal = None
738
-
739
- send1: ObjectSendStream[
740
- typing.Optional[typing.MutableMapping[str, typing.Any]]
741
- ]
742
- receive1: ObjectReceiveStream[
743
- typing.Optional[typing.MutableMapping[str, typing.Any]]
744
- ]
745
- send2: ObjectSendStream[typing.MutableMapping[str, typing.Any]]
746
- receive2: ObjectReceiveStream[typing.MutableMapping[str, typing.Any]]
747
- send1, receive1 = anyio.create_memory_object_stream(math.inf)
748
- send2, receive2 = anyio.create_memory_object_stream(math.inf)
749
- self.stream_send = StapledObjectStream(send1, receive1)
750
- self.stream_receive = StapledObjectStream(send2, receive2)
751
- self.task = portal.start_task_soon(self.lifespan)
752
- portal.call(self.wait_startup)
753
-
754
- @stack.callback
755
- def wait_shutdown() -> None:
756
- portal.call(self.wait_shutdown)
757
-
758
- self.exit_stack = stack.pop_all()
759
-
760
- return self
761
-
762
- def __exit__(self, *args: typing.Any) -> None:
763
- self.exit_stack.close()
764
-
765
- async def lifespan(self) -> None:
766
- scope = {"type": "lifespan", "state": self.app_state}
767
- try:
768
- await self.app(scope, self.stream_receive.receive, self.stream_send.send)
769
- finally:
770
- await self.stream_send.send(None)
771
-
772
- async def wait_startup(self) -> None:
773
- await self.stream_receive.send({"type": "lifespan.startup"})
774
-
775
- async def receive() -> typing.Any:
776
- message = await self.stream_send.receive()
777
- if message is None:
778
- self.task.result()
779
- return message
780
-
781
- message = await receive()
782
- assert message["type"] in (
783
- "lifespan.startup.complete",
784
- "lifespan.startup.failed",
785
- )
786
- if message["type"] == "lifespan.startup.failed":
787
- await receive()
788
-
789
- async def wait_shutdown(self) -> None:
790
- async def receive() -> typing.Any:
791
- message = await self.stream_send.receive()
792
- if message is None:
793
- self.task.result()
794
- return message
795
-
796
- async with self.stream_send:
797
- await self.stream_receive.send({"type": "lifespan.shutdown"})
798
- message = await receive()
799
- assert message["type"] in (
800
- "lifespan.shutdown.complete",
801
- "lifespan.shutdown.failed",
802
- )
803
- if message["type"] == "lifespan.shutdown.failed":
804
- await receive()