Skip to content

Commit 6006a54

Browse files
committed
Custom NoPolling implementations for return types
1 parent 0b5fc0c commit 6006a54

File tree

5 files changed

+97
-26
lines changed

5 files changed

+97
-26
lines changed

Diff for: sdk/keyvault/azure-keyvault-securitydomain/azure/keyvault/securitydomain/_internal/__init__.py

+12-1
Original file line numberDiff line numberDiff line change
@@ -6,27 +6,38 @@
66

77
from urllib.parse import urlparse
88

9-
from .async_polling import AsyncSecurityDomainDownloadPollingMethod, AsyncSecurityDomainUploadPollingMethod
9+
from .async_polling import (
10+
AsyncSecurityDomainDownloadNoPolling,
11+
AsyncSecurityDomainDownloadPollingMethod,
12+
AsyncSecurityDomainUploadNoPolling,
13+
AsyncSecurityDomainUploadPollingMethod,
14+
)
1015
from .challenge_auth_policy import ChallengeAuthPolicy
1116
from .http_challenge import HttpChallenge
1217
from . import http_challenge_cache
1318
from .polling import (
19+
SecurityDomainDownloadNoPolling,
1420
SecurityDomainDownloadPolling,
1521
SecurityDomainDownloadPollingMethod,
22+
SecurityDomainUploadNoPolling,
1623
SecurityDomainUploadPolling,
1724
SecurityDomainUploadPollingMethod,
1825
)
1926

2027
HttpChallengeCache = http_challenge_cache # to avoid aliasing pylint error (C4745)
2128

2229
__all__ = [
30+
"AsyncSecurityDomainDownloadNoPolling",
2331
"AsyncSecurityDomainDownloadPollingMethod",
32+
"AsyncSecurityDomainUploadNoPolling",
2433
"AsyncSecurityDomainUploadPollingMethod",
2534
"ChallengeAuthPolicy",
2635
"HttpChallenge",
2736
"HttpChallengeCache",
37+
"SecurityDomainDownloadNoPolling",
2838
"SecurityDomainDownloadPolling",
2939
"SecurityDomainDownloadPollingMethod",
40+
"SecurityDomainUploadNoPolling",
3041
"SecurityDomainUploadPolling",
3142
"SecurityDomainUploadPollingMethod",
3243
]

Diff for: sdk/keyvault/azure-keyvault-securitydomain/azure/keyvault/securitydomain/_internal/async_polling.py

+14-1
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,25 @@
44
# ------------------------------------
55
from azure.core.polling.async_base_polling import AsyncLROBasePolling
66

7-
from .polling import SecurityDomainDownloadPollingMethod, SecurityDomainUploadPollingMethod
7+
from .polling import (
8+
SecurityDomainDownloadPollingMethod,
9+
SecurityDomainDownloadNoPolling,
10+
SecurityDomainUploadPollingMethod,
11+
SecurityDomainUploadNoPolling,
12+
)
813

914

1015
class AsyncSecurityDomainDownloadPollingMethod(AsyncLROBasePolling, SecurityDomainDownloadPollingMethod):
1116
pass
1217

1318

19+
class AsyncSecurityDomainDownloadNoPolling(AsyncLROBasePolling, SecurityDomainDownloadNoPolling):
20+
pass
21+
22+
1423
class AsyncSecurityDomainUploadPollingMethod(AsyncLROBasePolling, SecurityDomainUploadPollingMethod):
1524
pass
25+
26+
27+
class AsyncSecurityDomainUploadNoPolling(AsyncLROBasePolling, SecurityDomainUploadNoPolling):
28+
pass

Diff for: sdk/keyvault/azure-keyvault-securitydomain/azure/keyvault/securitydomain/_internal/polling.py

+43-4
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616

1717
PollingReturnType_co = TypeVar("PollingReturnType_co", covariant=True)
1818

19-
# The correct success response should be "Succeeded", but this has already shipped.
20-
_FINISHED = frozenset(["success", "canceled", "failed"])
19+
# The correct success response should be "Succeeded", but this has already shipped. Handle "Succeeded" just in case.
20+
_FINISHED = frozenset(["succeeded", "success", "canceled", "failed"])
2121

2222

2323
def _finished(status):
@@ -71,6 +71,27 @@ def parse_resource(
7171
return None # type: ignore
7272

7373

74+
class NoPollingMixin(LROBasePolling):
75+
def finished(self) -> bool:
76+
"""Is this polling finished?
77+
78+
:rtype: bool
79+
:return: Whether this polling is finished
80+
"""
81+
return True
82+
83+
def status(self) -> str:
84+
"""Return the current status.
85+
86+
:rtype: str
87+
:return: The current status
88+
"""
89+
return "succeeded"
90+
91+
def result(self, *args, **kwargs):
92+
return self.resource()
93+
94+
7495
class SecurityDomainDownloadPolling(OperationResourcePolling):
7596
def __init__(self) -> None:
7697
self._polling_url = ""
@@ -121,7 +142,9 @@ def get_long_running_output(pipeline_response):
121142
response_headers["Azure-AsyncOperation"] = deserializer._deserialize( # pylint: disable=protected-access
122143
"str", response.headers.get("Azure-AsyncOperation")
123144
)
124-
response_headers["Retry-After"] = deserializer._deserialize("int", response.headers.get("Retry-After")) # pylint: disable=protected-access
145+
response_headers["Retry-After"] = deserializer._deserialize(
146+
"int", response.headers.get("Retry-After")
147+
) # pylint: disable=protected-access
125148

126149
return _deserialize(SecurityDomainObject, response.json())
127150

@@ -137,6 +160,10 @@ def resource(self) -> SecurityDomainObject:
137160
return cast(SecurityDomainObject, self.parse_resource(self._initial_response))
138161

139162

163+
class SecurityDomainDownloadNoPolling(SecurityDomainDownloadPollingMethod, NoPollingMixin):
164+
pass
165+
166+
140167
class SecurityDomainUploadPolling(SecurityDomainDownloadPolling):
141168
def set_initial_status(self, pipeline_response: "PipelineResponse") -> str:
142169
response: HttpResponse = pipeline_response.http_response
@@ -175,7 +202,9 @@ def get_long_running_output(pipeline_response):
175202
response_headers["Azure-AsyncOperation"] = deserializer._deserialize( # pylint: disable=protected-access
176203
"str", response.headers.get("Azure-AsyncOperation")
177204
)
178-
response_headers["Retry-After"] = deserializer._deserialize("int", response.headers.get("Retry-After")) # pylint: disable=protected-access
205+
response_headers["Retry-After"] = deserializer._deserialize(
206+
"int", response.headers.get("Retry-After")
207+
) # pylint: disable=protected-access
179208

180209
return _deserialize(SecurityDomainOperationStatus, response.json())
181210

@@ -188,3 +217,13 @@ def resource(self) -> SecurityDomainOperationStatus:
188217
:return: The built resource.
189218
"""
190219
return cast(SecurityDomainOperationStatus, self.parse_resource(self._pipeline_response))
220+
221+
222+
class SecurityDomainUploadNoPolling(SecurityDomainUploadPollingMethod, NoPollingMixin):
223+
def resource(self) -> SecurityDomainOperationStatus:
224+
"""Return the built resource.
225+
226+
:rtype: any
227+
:return: The built resource.
228+
"""
229+
return cast(SecurityDomainOperationStatus, self.parse_resource(self._initial_response))

Diff for: sdk/keyvault/azure-keyvault-securitydomain/azure/keyvault/securitydomain/_patch.py

+13-10
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,10 @@
2121
from ._client import KeyVaultClient
2222
from ._internal import (
2323
ChallengeAuthPolicy,
24+
SecurityDomainDownloadNoPolling,
2425
SecurityDomainDownloadPolling,
2526
SecurityDomainDownloadPollingMethod,
27+
SecurityDomainUploadNoPolling,
2628
SecurityDomainUploadPolling,
2729
SecurityDomainUploadPollingMethod,
2830
)
@@ -124,8 +126,7 @@ def begin_download(
124126
content_type: str = "application/json",
125127
polling: Optional[Literal[False]] = None,
126128
**kwargs: Any,
127-
) -> LROPoller[SecurityDomainObject]:
128-
...
129+
) -> LROPoller[SecurityDomainObject]: ...
129130

130131
@overload
131132
def begin_download(
@@ -135,8 +136,7 @@ def begin_download(
135136
content_type: str = "application/json",
136137
polling: Optional[Literal[False]] = None,
137138
**kwargs: Any,
138-
) -> LROPoller[SecurityDomainObject]:
139-
...
139+
) -> LROPoller[SecurityDomainObject]: ...
140140

141141
@overload
142142
def begin_download(
@@ -146,8 +146,7 @@ def begin_download(
146146
content_type: str = "application/json",
147147
polling: Optional[Literal[False]] = None,
148148
**kwargs: Any,
149-
) -> LROPoller[SecurityDomainObject]:
150-
...
149+
) -> LROPoller[SecurityDomainObject]: ...
151150

152151
@distributed_trace
153152
def begin_download(
@@ -178,8 +177,10 @@ def begin_download(
178177
:raises ~azure.core.exceptions.HttpResponseError:
179178
"""
180179
delay = kwargs.pop("polling_interval", self._config.polling_interval)
181-
polling_method = False if polling is False else SecurityDomainDownloadPollingMethod(
182-
lro_algorithms=[SecurityDomainDownloadPolling()], timeout=delay
180+
polling_method = (
181+
SecurityDomainDownloadNoPolling()
182+
if polling is False
183+
else SecurityDomainDownloadPollingMethod(lro_algorithms=[SecurityDomainDownloadPolling()], timeout=delay)
183184
)
184185
return super().begin_download( # type: ignore[return-value]
185186
certificate_info_object,
@@ -215,8 +216,10 @@ def begin_upload(
215216
:raises ~azure.core.exceptions.HttpResponseError:
216217
"""
217218
delay = kwargs.pop("polling_interval", self._config.polling_interval)
218-
polling_method = False if polling is False else SecurityDomainUploadPollingMethod(
219-
lro_algorithms=[SecurityDomainUploadPolling()], timeout=delay
219+
polling_method = (
220+
SecurityDomainUploadNoPolling()
221+
if polling is False
222+
else SecurityDomainUploadPollingMethod(lro_algorithms=[SecurityDomainUploadPolling()], timeout=delay)
220223
)
221224
return super().begin_upload(
222225
security_domain,

Diff for: sdk/keyvault/azure-keyvault-securitydomain/azure/keyvault/securitydomain/aio/_patch.py

+15-10
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@
1818
from ._client import KeyVaultClient
1919
from .._internal import (
2020
AsyncChallengeAuthPolicy,
21+
AsyncSecurityDomainDownloadNoPolling,
2122
AsyncSecurityDomainDownloadPollingMethod,
23+
AsyncSecurityDomainUploadNoPolling,
2224
AsyncSecurityDomainUploadPollingMethod,
2325
SecurityDomainDownloadPolling,
2426
SecurityDomainUploadPolling,
@@ -85,8 +87,7 @@ async def begin_download(
8587
content_type: str = "application/json",
8688
polling: Optional[Literal[False]] = None,
8789
**kwargs: Any,
88-
) -> AsyncLROPoller[SecurityDomainObject]:
89-
...
90+
) -> AsyncLROPoller[SecurityDomainObject]: ...
9091

9192
@overload
9293
async def begin_download(
@@ -96,8 +97,7 @@ async def begin_download(
9697
content_type: str = "application/json",
9798
polling: Optional[Literal[False]] = None,
9899
**kwargs: Any,
99-
) -> AsyncLROPoller[SecurityDomainObject]:
100-
...
100+
) -> AsyncLROPoller[SecurityDomainObject]: ...
101101

102102
@overload
103103
async def begin_download(
@@ -107,8 +107,7 @@ async def begin_download(
107107
content_type: str = "application/json",
108108
polling: Optional[Literal[False]] = None,
109109
**kwargs: Any,
110-
) -> AsyncLROPoller[SecurityDomainObject]:
111-
...
110+
) -> AsyncLROPoller[SecurityDomainObject]: ...
112111

113112
@distributed_trace_async
114113
async def begin_download(
@@ -139,8 +138,12 @@ async def begin_download(
139138
:raises ~azure.core.exceptions.HttpResponseError:
140139
"""
141140
delay = kwargs.pop("polling_interval", self._config.polling_interval)
142-
polling_method = False if polling is False else AsyncSecurityDomainDownloadPollingMethod(
143-
lro_algorithms=[SecurityDomainDownloadPolling()], timeout=delay
141+
polling_method = (
142+
AsyncSecurityDomainDownloadNoPolling()
143+
if polling is False
144+
else AsyncSecurityDomainDownloadPollingMethod(
145+
lro_algorithms=[SecurityDomainDownloadPolling()], timeout=delay
146+
)
144147
)
145148
return await super().begin_download( # type: ignore[return-value]
146149
certificate_info_object,
@@ -176,8 +179,10 @@ async def begin_upload(
176179
:raises ~azure.core.exceptions.HttpResponseError:
177180
"""
178181
delay = kwargs.pop("polling_interval", self._config.polling_interval)
179-
polling_method = False if polling is False else AsyncSecurityDomainUploadPollingMethod(
180-
lro_algorithms=[SecurityDomainUploadPolling()], timeout=delay
182+
polling_method = (
183+
AsyncSecurityDomainUploadNoPolling()
184+
if polling is False
185+
else AsyncSecurityDomainUploadPollingMethod(lro_algorithms=[SecurityDomainUploadPolling()], timeout=delay)
181186
)
182187
return await super().begin_upload(
183188
security_domain,

0 commit comments

Comments
 (0)