Skip to content

Commit 82ec61c

Browse files
BenRKarlBenRKarl
committed
Add type annotations to interceptors/response_wrappers.py
Change-Id: I174214176c1a2739be92436161ceea1ebbbd4d5b
1 parent 526a2d3 commit 82ec61c

File tree

1 file changed

+74
-58
lines changed

1 file changed

+74
-58
lines changed

‎google/ads/googleads/interceptors/response_wrappers.py‎

Lines changed: 74 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -14,169 +14,185 @@
1414
"""Wrapper classes used to modify the behavior of response objects."""
1515

1616
import grpc
17-
import time
17+
from types import SimpleNamespace, TracebackType
18+
from typing import Any, Callable, Iterator, Optional, Union
1819

1920
from google.ads.googleads import util
20-
from types import SimpleNamespace
21+
from google.ads.googleads.interceptors import MetadataType
22+
from google.protobuf.message import Message as ProtobufMessageType
2123

2224

2325
class _UnaryStreamWrapper(grpc.Call, grpc.Future):
24-
def __init__(self, underlay_call, failure_handler, use_proto_plus=False):
26+
def __init__(
27+
self,
28+
underlay_call: Union[grpc.Call, grpc.Future],
29+
failure_handler: Callable[[Union[grpc.Call, grpc.Future]], grpc.Call],
30+
use_proto_plus: bool = False,
31+
):
2532
super().__init__()
26-
self._underlay_call = underlay_call
27-
self._failure_handler = failure_handler
28-
self._exception = None
29-
self._use_proto_plus = use_proto_plus
30-
self._cache = SimpleNamespace(**{"initial_response_object": None})
31-
32-
def initial_metadata(self):
33+
self._underlay_call: Union[grpc.Call, grpc.Future] = underlay_call
34+
self._failure_handler: Callable[[Union[grpc.Call, grpc.Future]], grpc.Call] = (
35+
failure_handler
36+
)
37+
self._exception: Optional[grpc.RpcError] = None
38+
self._use_proto_plus: bool = use_proto_plus
39+
self._cache: SimpleNamespace = SimpleNamespace(
40+
**{"initial_response_object": None}
41+
)
42+
43+
def initial_metadata(self) -> MetadataType:
3344
return self._underlay_call.initial_metadata()
3445

35-
def trailing_metadata(self):
46+
def trailing_metadata(self) -> MetadataType:
3647
return self._underlay_call.initial_metadata()
3748

38-
def code(self):
49+
def code(self) -> grpc.StatusCode:
3950
return self._underlay_call.code()
4051

41-
def details(self):
52+
def details(self) -> str:
4253
return self._underlay_call.details()
4354

44-
def debug_error_string(self):
55+
def debug_error_string(self) -> str:
4556
return self._underlay_call.debug_error_string()
4657

47-
def cancelled(self):
58+
def cancelled(self) -> bool:
4859
return self._underlay_call.cancelled()
4960

50-
def running(self):
61+
def running(self) -> bool:
5162
return self._underlay_call.running()
5263

53-
def done(self):
64+
def done(self) -> bool:
5465
return self._underlay_call.done()
5566

56-
def result(self, timeout=None):
67+
def result(self, timeout: Optional[float] = None) -> ProtobufMessageType:
5768
return self._underlay_call.result(timeout=timeout)
5869

59-
def exception(self, timeout=None):
70+
def exception(self, timeout: Optional[float] = None) -> Optional[grpc.RpcError]:
6071
if self._exception:
6172
return self._exception
6273
else:
6374
return self._underlay_call.exception(timeout=timeout)
6475

65-
def traceback(self, timeout=None):
76+
def traceback(self, timeout: Optional[float] = None) -> Optional[TracebackType]:
6677
return self._underlay_call.traceback(timeout=timeout)
6778

68-
def add_done_callback(self, fn):
79+
def add_done_callback(self, fn: Callable[[grpc.Future], Any]) -> None:
6980
return self._underlay_call.add_done_callback(fn)
7081

71-
def add_callback(self, callback):
82+
def add_callback(self, callback: Callable[[grpc.Future], Any]) -> None:
7283
return self._underlay_call.add_callback(callback)
7384

74-
def is_active(self):
85+
def is_active(self) -> bool:
7586
return self._underlay_call.is_active()
7687

77-
def time_remaining(self):
88+
def time_remaining(self) -> Optional[float]:
7889
return self._underlay_call.time_remaining()
7990

80-
def cancel(self):
91+
def cancel(self) -> bool:
8192
return self._underlay_call.cancel()
8293

83-
def __iter__(self):
94+
def __iter__(self) -> Iterator[ProtobufMessageType]:
8495
return self
8596

86-
def __next__(self):
97+
def __next__(self) -> ProtobufMessageType:
8798
try:
88-
message = next(self._underlay_call)
99+
message: ProtobufMessageType = next(self._underlay_call)
89100
# Store only the first streaming response object in _cache.initial_response_object.
90101
# Each streaming response object captures 10,000 rows.
91102
# The default log character limit is 5,000, so caching multiple
92103
# streaming response objects does not make sense in most cases,
93104
# as only [part of] 1 will get logged.
94-
if self._cache.initial_response_object == None:
105+
if self._cache.initial_response_object is None:
95106
self._cache.initial_response_object = message
96-
if self._use_proto_plus == True:
107+
if self._use_proto_plus is True:
97108
# By default this message is wrapped by proto-plus
98109
return message
99110
else:
100111
return util.convert_proto_plus_to_protobuf(message)
101112
except StopIteration:
102113
raise
103-
except Exception:
114+
except Exception as e:
104115
try:
105116
self._failure_handler(self._underlay_call)
106-
except Exception as e:
117+
except Exception as e_inner:
107118
self._exception = e
108-
raise e
119+
raise e_inner
109120

110-
def get_cache(self):
121+
def get_cache(self) -> SimpleNamespace:
111122
return self._cache
112123

113124

114125
class _UnaryUnaryWrapper(grpc.Call, grpc.Future):
115-
def __init__(self, underlay_call, use_proto_plus=False):
126+
def __init__(
127+
self,
128+
underlay_call: Union[grpc.Call, grpc.Future],
129+
use_proto_plus: bool = False,
130+
):
116131
super().__init__()
117-
self._underlay_call = underlay_call
118-
self._use_proto_plus = use_proto_plus
132+
self._underlay_call: Union[grpc.Call, grpc.Future] = underlay_call
133+
self._use_proto_plus: bool = use_proto_plus
134+
self._exception: Optional[grpc.RpcError] = None
119135

120-
def initial_metadata(self):
136+
def initial_metadata(self) -> MetadataType:
121137
return self._underlay_call.initial_metadata()
122138

123-
def trailing_metadata(self):
139+
def trailing_metadata(self) -> MetadataType:
124140
return self._underlay_call.initial_metadata()
125141

126-
def code(self):
142+
def code(self) -> grpc.StatusCode:
127143
return self._underlay_call.code()
128144

129-
def details(self):
145+
def details(self) -> str:
130146
return self._underlay_call.details()
131147

132-
def debug_error_string(self):
148+
def debug_error_string(self) -> str:
133149
return self._underlay_call.debug_error_string()
134150

135-
def cancelled(self):
151+
def cancelled(self) -> bool:
136152
return self._underlay_call.cancelled()
137153

138-
def running(self):
154+
def running(self) -> bool:
139155
return self._underlay_call.running()
140156

141-
def done(self):
157+
def done(self) -> bool:
142158
return self._underlay_call.done()
143159

144-
def result(self, timeout=None):
145-
message = self._underlay_call.result()
146-
if self._use_proto_plus == True:
160+
def result(self, timeout: Optional[float] = None) -> ProtobufMessageType:
161+
message: ProtobufMessageType = self._underlay_call.result()
162+
if self._use_proto_plus is True:
147163
return message
148164
else:
149165
return util.convert_proto_plus_to_protobuf(message)
150166

151-
def exception(self, timeout=None):
167+
def exception(self, timeout: Optional[float] = None) -> Optional[grpc.RpcError]:
152168
if self._exception:
153169
return self._exception
154170
else:
155171
return self._underlay_call.exception(timeout=timeout)
156172

157-
def traceback(self, timeout=None):
173+
def traceback(self, timeout: Optional[float] = None) -> Optional[TracebackType]:
158174
return self._underlay_call.traceback(timeout=timeout)
159175

160-
def add_done_callback(self, fn):
176+
def add_done_callback(self, fn: Callable[[grpc.Future], Any]) -> None:
161177
return self._underlay_call.add_done_callback(fn)
162178

163-
def add_callback(self, callback):
179+
def add_callback(self, callback: Callable[[grpc.Future], Any]) -> None:
164180
return self._underlay_call.add_callback(callback)
165181

166-
def is_active(self):
182+
def is_active(self) -> bool:
167183
return self._underlay_call.is_active()
168184

169-
def time_remaining(self):
185+
def time_remaining(self) -> Optional[float]:
170186
return self._underlay_call.time_remaining()
171187

172-
def cancel(self):
188+
def cancel(self) -> bool:
173189
return self._underlay_call.cancel()
174190

175-
def __iter__(self):
176-
if self._use_proto_plus == True:
191+
def __iter__(self) -> Iterator[ProtobufMessageType]:
192+
if self._use_proto_plus is True:
177193
return self
178194
else:
179195
return util.convert_proto_plus_to_protobuf(self)
180196

181-
def __next__(self):
197+
def __next__(self) -> ProtobufMessageType:
182198
return next(self._underlay_call)

0 commit comments

Comments
 (0)