Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit a04211b

Browse files
eli-bleli-bl
committedNov 15, 2024·
fix template
1 parent cf30e48 commit a04211b

File tree

2 files changed

+158
-60
lines changed
  • end_to_end_tests/golden-record-dataclasses/my_dataclasses_api_client
  • openapi_python_client/templates

2 files changed

+158
-60
lines changed
 

‎end_to_end_tests/golden-record-dataclasses/my_dataclasses_api_client/client.py

Lines changed: 69 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import copy
1+
from copy import copy
22
from dataclasses import dataclass, field, replace
33
from typing import Any, Dict, Optional
44

@@ -44,8 +44,9 @@ class Client:
4444
"""
4545

4646
_base_url: str
47-
raise_on_unexpected_status: bool
4847
_httpx_config: _HttpxConfig
48+
_client: Optional[httpx.Client]
49+
_async_client: Optional[httpx.AsyncClient]
4950

5051
def __init__(
5152
self,
@@ -61,26 +62,42 @@ def __init__(
6162
self.raise_on_unexpected_status = raise_on_unexpected_status
6263
self._base_url = base_url
6364
self._httpx_config = _HttpxConfig(cookies, headers, timeout, verify_ssl, follow_redirects, httpx_args)
65+
self._client = None
66+
self._async_client = None
6467

6568
def _with_httpx_config(self, httpx_config: _HttpxConfig) -> "Client":
66-
ret = copy.copy(self)
69+
ret = copy(self)
6770
ret._httpx_config = httpx_config
71+
ret._client = None
72+
ret._async_client = None
6873
return ret
6974

7075
def with_headers(self, headers: Dict[str, str]) -> "Client":
7176
"""Get a new client matching this one with additional headers"""
77+
if self._client is not None:
78+
self._client.headers.update(headers)
79+
if self._async_client is not None:
80+
self._async_client.headers.update(headers)
7281
return self._with_httpx_config(
7382
replace(self._httpx_config, headers={**self._httpx_config.headers, **headers}),
7483
)
7584

7685
def with_cookies(self, cookies: Dict[str, str]) -> "Client":
7786
"""Get a new client matching this one with additional cookies"""
87+
if self._client is not None:
88+
self._client.cookies.update(cookies)
89+
if self._async_client is not None:
90+
self._async_client.cookies.update(cookies)
7891
return self._with_httpx_config(
7992
replace(self._httpx_config, cookies={**self._httpx_config.cookies, **cookies}),
8093
)
8194

8295
def with_timeout(self, timeout: httpx.Timeout) -> "Client":
8396
"""Get a new client matching this one with a new timeout (in seconds)"""
97+
if self._client is not None:
98+
self._client.timeout = timeout
99+
if self._async_client is not None:
100+
self._async_client.timeout = timeout
84101
return self._with_httpx_config(replace(self._httpx_config, timeout=timeout))
85102

86103
def set_httpx_client(self, client: httpx.Client) -> "Client":
@@ -178,31 +195,70 @@ class AuthenticatedClient:
178195
"""
179196

180197
_base_url: str
181-
raise_on_unexpected_status: bool
182198
token: str
183199
prefix: str = "Bearer"
184200
auth_header_name: str = "Authorization"
185201
_httpx_config: _HttpxConfig
202+
_client: Optional[httpx.Client]
203+
_async_client: Optional[httpx.AsyncClient]
204+
205+
def __init__(
206+
self,
207+
base_url: str,
208+
token: str,
209+
prefix: str = "Bearer",
210+
auth_header_name: str = "Authorization",
211+
cookies: Dict[str, str] = {},
212+
headers: Dict[str, str] = {},
213+
timeout: Optional[httpx.Timeout] = None,
214+
verify_ssl: bool = True,
215+
follow_redirects: bool = False,
216+
httpx_args: Dict[str, Any] = {},
217+
raise_on_unexpected_status: bool = False,
218+
) -> None:
219+
self.raise_on_unexpected_status = raise_on_unexpected_status
220+
self._base_url = base_url
221+
self._httpx_config = _HttpxConfig(cookies, headers, timeout, verify_ssl, follow_redirects, httpx_args)
222+
self._client = None
223+
self._async_client = None
224+
225+
self.token = token
226+
self.prefix = prefix
227+
self.auth_header_name = auth_header_name
186228

187229
def _with_httpx_config(self, httpx_config: _HttpxConfig) -> "AuthenticatedClient":
188-
ret = copy.copy(self)
230+
ret = copy(self)
189231
ret._httpx_config = httpx_config
232+
ret._client = None
233+
ret._async_client = None
190234
return ret
191235

192236
def with_headers(self, headers: Dict[str, str]) -> "AuthenticatedClient":
193237
"""Get a new client matching this one with additional headers"""
238+
if self._client is not None:
239+
self._client.headers.update(headers)
240+
if self._async_client is not None:
241+
self._async_client.headers.update(headers)
194242
return self._with_httpx_config(
195243
replace(self._httpx_config, headers={**self._httpx_config.headers, **headers}),
196244
)
197245

198246
def with_cookies(self, cookies: Dict[str, str]) -> "AuthenticatedClient":
199247
"""Get a new client matching this one with additional cookies"""
248+
if self._client is not None:
249+
self._client.cookies.update(cookies)
250+
if self._async_client is not None:
251+
self._async_client.cookies.update(cookies)
200252
return self._with_httpx_config(
201253
replace(self._httpx_config, cookies={**self._httpx_config.cookies, **cookies}),
202254
)
203255

204256
def with_timeout(self, timeout: httpx.Timeout) -> "AuthenticatedClient":
205257
"""Get a new client matching this one with a new timeout (in seconds)"""
258+
if self._client is not None:
259+
self._client.timeout = timeout
260+
if self._async_client is not None:
261+
self._async_client.timeout = timeout
206262
return self._with_httpx_config(replace(self._httpx_config, timeout=timeout))
207263

208264
def set_httpx_client(self, client: httpx.Client) -> "AuthenticatedClient":
@@ -216,11 +272,13 @@ def set_httpx_client(self, client: httpx.Client) -> "AuthenticatedClient":
216272
def get_httpx_client(self) -> httpx.Client:
217273
"""Get the underlying httpx.Client, constructing a new one if not previously set"""
218274
if self._client is None:
219-
self._headers[self.auth_header_name] = f"{self.prefix} {self.token}" if self.prefix else self.token
220275
self._client = httpx.Client(
221276
base_url=self._base_url,
222277
cookies=self._httpx_config.cookies,
223-
headers=self._httpx_config.headers,
278+
headers={
279+
self.auth_header_name: (f"{self.prefix} {self.token}" if self.prefix else self.token),
280+
**self._httpx_config.headers,
281+
},
224282
timeout=self._httpx_config.timeout,
225283
verify=self._httpx_config.verify_ssl,
226284
follow_redirects=self._httpx_config.follow_redirects,
@@ -248,11 +306,13 @@ def set_async_httpx_client(self, async_client: httpx.AsyncClient) -> "Authentica
248306
def get_async_httpx_client(self) -> httpx.AsyncClient:
249307
"""Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
250308
if self._async_client is None:
251-
self._headers[self.auth_header_name] = f"{self.prefix} {self.token}" if self.prefix else self.token
252309
self._async_client = httpx.AsyncClient(
253310
base_url=self._httpx_config.base_url,
254311
cookies=self._httpx_config.cookies,
255-
headers=self._httpx_config.headers,
312+
headers={
313+
self.auth_header_name: (f"{self.prefix} {self.token}" if self.prefix else self.token),
314+
**self._httpx_config.headers,
315+
},
256316
timeout=self._httpx_config.timeout,
257317
verify=self._httpx_config.verify_ssl,
258318
follow_redirects=self._httpx_config.follow_redirects,

‎openapi_python_client/templates/client.py.jinja

Lines changed: 89 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@ from typing import Any, Dict, Union, Optional
44
import httpx
55

66
{% if config.use_dataclasses %}
7-
import copy
7+
from copy import copy
88
from dataclasses import dataclass, field, replace
9+
{% set class_header = "" %}
910
{% else %}
1011
from attrs import define, field, evolve
12+
{% set class_header = "@define" %}
1113
{% endif %}
1214

1315
{% if config.use_dataclasses %}
@@ -19,27 +21,13 @@ class _HttpxConfig:
1921
verify_ssl: bool = True
2022
follow_redirects: bool = False
2123
httpx_args: Dict[str, str] = field(default_factory=dict)
22-
23-
{% macro init_args() %}
24-
self,
25-
base_url: str,
26-
cookies: Dict[str, str] = {},
27-
headers: Dict[str, str] = {},
28-
timeout: Optional[httpx.Timeout] = None,
29-
verify_ssl: bool = True,
30-
follow_redirects: bool = False,
31-
httpx_args: Dict[str, Any] = {},
32-
raise_on_unexpected_status: bool = False,
33-
{% endmacro %}
3424
{% endif %}
3525

3626
{% macro httpx_client_arg(name) -%}
3727
{% if config.use_dataclasses %}self._httpx_config.{{ name }}{% else %}self._{{ name }}{% endif %}
3828
{%- endmacro %}
3929

40-
{% if not config.use_dataclasses %}
41-
@define
42-
{% endif %}
30+
{{ class_header }}
4331
class Client:
4432
"""A class for keeping track of data related to the API
4533

@@ -72,11 +60,12 @@ class Client:
7260
{% macro attributes(extra_attrs = []) %}
7361
{% if config.use_dataclasses %}
7462
_base_url: str
75-
raise_on_unexpected_status: bool
7663
{% for extra in extra_attrs %}
7764
{{ extra }}
7865
{% endfor %}
7966
_httpx_config: _HttpxConfig
67+
_client: Optional[httpx.Client]
68+
_async_client: Optional[httpx.AsyncClient]
8069
{% else %}
8170
raise_on_unexpected_status: bool = field(default=False, kw_only=True)
8271
_base_url: str = field(alias="base_url")
@@ -95,69 +84,92 @@ class Client:
9584
{% endfor %}
9685
{% endif %}
9786
{% endif %}
98-
{% endmacro %}{{ attributes() }}
87+
{% endmacro %}
9988

100-
{% if config.use_dataclasses %}
89+
{{ attributes() }}
90+
91+
{% macro explicit_initializer(extra_attrs = []) %}
10192
def __init__(
102-
{{ init_args() }}
93+
self,
94+
base_url: str,
95+
{% for attr in extra_attrs %}
96+
{{ attr }},
97+
{% endfor %}
98+
cookies: Dict[str, str] = {},
99+
headers: Dict[str, str] = {},
100+
timeout: Optional[httpx.Timeout] = None,
101+
verify_ssl: bool = True,
102+
follow_redirects: bool = False,
103+
httpx_args: Dict[str, Any] = {},
104+
raise_on_unexpected_status: bool = False,
103105
) -> None:
104106
self.raise_on_unexpected_status = raise_on_unexpected_status
105107
self._base_url = base_url
106108
self._httpx_config = _HttpxConfig(
107109
cookies, headers, timeout, verify_ssl, follow_redirects, httpx_args
108110
)
111+
self._client = None
112+
self._async_client = None
113+
{% endmacro %}
114+
115+
{% if config.use_dataclasses %}
116+
{{ explicit_initializer() }}
109117
{% endif %}
110118

111119
{% macro builders(self) %}
112120
{% if config.use_dataclasses %}
113121
def _with_httpx_config(self, httpx_config: _HttpxConfig) -> "{{ self }}":
114-
ret = copy.copy(self)
122+
ret = copy(self)
115123
ret._httpx_config = httpx_config
124+
ret._client = None
125+
ret._async_client = None
116126
return ret
127+
{% endif %}
117128

118-
def with_headers(self, headers: Dict[str, str]) -> "{{ self }}":
119-
"""Get a new client matching this one with additional headers"""
120-
return self._with_httpx_config(
121-
replace(self._httpx_config, headers={**self._httpx_config.headers, **headers}),
122-
)
123-
124-
def with_cookies(self, cookies: Dict[str, str]) -> "{{ self }}":
125-
"""Get a new client matching this one with additional cookies"""
126-
return self._with_httpx_config(
127-
replace(self._httpx_config, cookies={**self._httpx_config.cookies, **cookies}),
128-
)
129-
130-
def with_timeout(self, timeout: httpx.Timeout) -> "{{ self }}":
131-
"""Get a new client matching this one with a new timeout (in seconds)"""
132-
return self._with_httpx_config(replace(self._httpx_config, timeout=timeout))
133-
{% else %}
134129
def with_headers(self, headers: Dict[str, str]) -> "{{ self }}":
135130
"""Get a new client matching this one with additional headers"""
136131
if self._client is not None:
137132
self._client.headers.update(headers)
138133
if self._async_client is not None:
139134
self._async_client.headers.update(headers)
135+
{% if config.use_dataclasses %}
136+
return self._with_httpx_config(
137+
replace(self._httpx_config, headers={**self._httpx_config.headers, **headers}),
138+
)
139+
{% else %}
140140
return evolve(self, headers={**self._headers, **headers})
141+
{% endif %}
141142

142143
def with_cookies(self, cookies: Dict[str, str]) -> "{{ self }}":
143144
"""Get a new client matching this one with additional cookies"""
144145
if self._client is not None:
145146
self._client.cookies.update(cookies)
146147
if self._async_client is not None:
147148
self._async_client.cookies.update(cookies)
149+
{% if config.use_dataclasses %}
150+
return self._with_httpx_config(
151+
replace(self._httpx_config, cookies={**self._httpx_config.cookies, **cookies}),
152+
)
153+
{% else %}
148154
return evolve(self, cookies={**self._cookies, **cookies})
155+
{% endif %}
149156

150157
def with_timeout(self, timeout: httpx.Timeout) -> "{{ self }}":
151158
"""Get a new client matching this one with a new timeout (in seconds)"""
152159
if self._client is not None:
153160
self._client.timeout = timeout
154161
if self._async_client is not None:
155162
self._async_client.timeout = timeout
163+
{% if config.use_dataclasses %}
164+
return self._with_httpx_config(replace(self._httpx_config, timeout=timeout))
165+
{% else %}
156166
return evolve(self, timeout=timeout)
157-
{% endif %}
158-
{% endmacro %}{{ builders("Client") }}
167+
{% endif %}
168+
{% endmacro %}
169+
170+
{{ builders("Client") }}
159171

160-
{% macro httpx_stuff(name, custom_constructor=None) %}
172+
{% macro httpx_stuff(name, headers_expr, update_headers=None) %}
161173
def set_httpx_client(self, client: httpx.Client) -> "{{ name }}":
162174
"""Manually set the underlying httpx.Client
163175

@@ -169,13 +181,13 @@ class Client:
169181
def get_httpx_client(self) -> httpx.Client:
170182
"""Get the underlying httpx.Client, constructing a new one if not previously set"""
171183
if self._client is None:
172-
{% if custom_constructor %}
173-
{{ custom_constructor | indent(12) }}
184+
{% if update_headers %}
185+
{{ update_headers | indent(12) }}
174186
{% endif %}
175187
self._client = httpx.Client(
176188
base_url=self._base_url,
177189
cookies={{ httpx_client_arg("cookies") }},
178-
headers={{ httpx_client_arg("headers") }},
190+
headers={{ headers_expr }},
179191
timeout={{ httpx_client_arg("timeout") }},
180192
verify={{ httpx_client_arg("verify_ssl") }},
181193
follow_redirects={{ httpx_client_arg("follow_redirects") }},
@@ -203,13 +215,13 @@ class Client:
203215
def get_async_httpx_client(self) -> httpx.AsyncClient:
204216
"""Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
205217
if self._async_client is None:
206-
{% if custom_constructor %}
207-
{{ custom_constructor | indent(12) }}
218+
{% if update_headers %}
219+
{{ update_headers | indent(12) }}
208220
{% endif %}
209221
self._async_client = httpx.AsyncClient(
210222
base_url={{ httpx_client_arg("base_url") }},
211223
cookies={{ httpx_client_arg("cookies") }},
212-
headers={{ httpx_client_arg("headers") }},
224+
headers={{ headers_expr or httpx_client_arg("headers") }},
213225
timeout={{ httpx_client_arg("timeout") }},
214226
verify={{ httpx_client_arg("verify_ssl") }},
215227
follow_redirects={{ httpx_client_arg("follow_redirects") }},
@@ -225,11 +237,15 @@ class Client:
225237
async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
226238
"""Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
227239
await self.get_async_httpx_client().__aexit__(*args, **kwargs)
228-
{% endmacro %}{{ httpx_stuff("Client") }}
240+
{% endmacro %}
229241

230-
{% if not config.use_dataclasses %}
231-
@define
242+
{% if config.use_dataclasses %}
243+
{{ httpx_stuff("Client", headers_expr="self._httpx_config.headers") }}
244+
{% else %}
245+
{{ httpx_stuff("Client", headers_expr="self._headers") }}
232246
{% endif %}
247+
248+
{{ class_header }}
233249
class AuthenticatedClient:
234250
"""A Client which has been authenticated for use on secured endpoints
235251

@@ -244,7 +260,29 @@ class AuthenticatedClient:
244260
auth_header_name: The name of the Authorization header
245261
"""
246262

247-
{{ attributes(["token: str", 'prefix: str = "Bearer"', 'auth_header_name: str = "Authorization"']) }}
263+
{% set extra_auth_client_args = ["token: str", 'prefix: str = "Bearer"', 'auth_header_name: str = "Authorization"'] %}
264+
{{ attributes(extra_auth_client_args) }}
265+
266+
{% if config.use_dataclasses %}
267+
{{ explicit_initializer(extra_auth_client_args) }}
268+
self.token = token
269+
self.prefix = prefix
270+
self.auth_header_name = auth_header_name
271+
{% endif %}
248272

249273
{{ builders("AuthenticatedClient") }}
250-
{{ httpx_stuff("AuthenticatedClient", "self._headers[self.auth_header_name] = f\"{self.prefix} {self.token}\" if self.prefix else self.token") }}
274+
{% set token_string_expr %} f"{self.prefix} {self.token}" if self.prefix else self.token {% endset %}
275+
{% if config.use_dataclasses %}
276+
{% set headers_expr %} { self.auth_header_name: ({{ token_string_expr }}), **self._httpx_config.headers } {% endset %}
277+
{{ httpx_stuff(
278+
"AuthenticatedClient",
279+
headers_expr=headers_expr,
280+
) }}
281+
{% else %}
282+
{% set update_headers %}self._headers[self.auth_header_name] = {{ token_string_expr }}{% endset %}
283+
{{ httpx_stuff(
284+
"AuthenticatedClient",
285+
headers_expr="self._headers",
286+
update_headers=update_headers,
287+
) }}
288+
{% endif %}

0 commit comments

Comments
 (0)
Please sign in to comment.