20
20
import weakref
21
21
from datetime import date , datetime
22
22
from functools import wraps
23
+ from typing import (
24
+ TYPE_CHECKING ,
25
+ Any ,
26
+ Callable ,
27
+ Collection ,
28
+ Dict ,
29
+ List ,
30
+ Optional ,
31
+ Tuple ,
32
+ TypeVar ,
33
+ Union ,
34
+ )
23
35
24
36
from ..compat import quote , string_types , to_bytes , to_str , unquote , urlparse
37
+ from ..serializer import Serializer
38
+ from ..transport import Transport
39
+
40
+ if TYPE_CHECKING :
41
+ from ..client import Elasticsearch
25
42
26
43
# parts of URL to be omitted
27
- SKIP_IN_PATH = (None , "" , b"" , [], ())
44
+ SKIP_IN_PATH : Collection [ Any ] = (None , "" , b"" , [], ())
28
45
29
46
30
- def _normalize_hosts (hosts ):
47
+ def _normalize_hosts (
48
+ hosts : Optional [Union [str , Collection [Union [str , Dict [str , Any ]]]]]
49
+ ) -> List [Dict [str , Any ]]:
31
50
"""
32
51
Helper function to transform hosts argument to
33
52
:class:`~elasticsearch.Elasticsearch` to a list of dicts.
@@ -40,15 +59,15 @@ def _normalize_hosts(hosts):
40
59
if isinstance (hosts , string_types ):
41
60
hosts = [hosts ]
42
61
43
- out = []
62
+ out : List [ Dict [ str , Any ]] = []
44
63
# normalize hosts to dicts
45
64
for host in hosts :
46
65
if isinstance (host , string_types ):
47
66
if "://" not in host :
48
67
host = f"//{ host } "
49
68
50
69
parsed_url = urlparse (host )
51
- h = {"host" : parsed_url .hostname }
70
+ h : Dict [ str , Any ] = {"host" : parsed_url .hostname }
52
71
53
72
if parsed_url .port :
54
73
h ["port" ] = parsed_url .port
@@ -59,20 +78,20 @@ def _normalize_hosts(hosts):
59
78
60
79
if parsed_url .username or parsed_url .password :
61
80
h ["http_auth" ] = "{}:{}" .format (
62
- unquote (parsed_url .username ),
63
- unquote (parsed_url .password ),
81
+ unquote (parsed_url .username or "" ),
82
+ unquote (parsed_url .password or "" ),
64
83
)
65
84
66
85
if parsed_url .path and parsed_url .path != "/" :
67
86
h ["url_prefix" ] = parsed_url .path
68
87
69
88
out .append (h )
70
89
else :
71
- out .append (host )
90
+ out .append (host ) # type: ignore
72
91
return out
73
92
74
93
75
- def _escape (value ) :
94
+ def _escape (value : Any ) -> Union [ str , bytes ] :
76
95
"""
77
96
Escape a single value of a URL string or a query parameter. If it is a list
78
97
or tuple, turn it into a comma-separated string first.
@@ -96,11 +115,11 @@ def _escape(value):
96
115
97
116
# encode strings to utf-8
98
117
if not isinstance (value , str ):
99
- value = str (value )
118
+ return str (value ). encode ( "utf-8" )
100
119
return value .encode ("utf-8" )
101
120
102
121
103
- def _make_path (* parts ) :
122
+ def _make_path (* parts : Any ) -> str :
104
123
"""
105
124
Create a URL string from parts, omit all `None` values and empty strings.
106
125
Convert lists and tuples to comma separated values.
@@ -115,18 +134,27 @@ def _make_path(*parts):
115
134
116
135
117
136
# parameters that apply to all methods
118
- GLOBAL_PARAMS = ("pretty" , "human" , "error_trace" , "format" , "filter_path" )
119
-
120
-
121
- def query_params (* es_query_params ):
137
+ GLOBAL_PARAMS : Tuple [str , ...] = (
138
+ "pretty" ,
139
+ "human" ,
140
+ "error_trace" ,
141
+ "format" ,
142
+ "filter_path" ,
143
+ )
144
+ T = TypeVar ("T" )
145
+
146
+
147
+ def query_params (
148
+ * es_query_params : str ,
149
+ ) -> Callable [[Callable [..., T ]], Callable [..., T ]]:
122
150
"""
123
151
Decorator that pops all accepted parameters from method's kwargs and puts
124
152
them in the params argument.
125
153
"""
126
154
127
- def _wrapper (func ) :
155
+ def _wrapper (func : Any ) -> Any :
128
156
@wraps (func )
129
- def _wrapped (* args , ** kwargs ) :
157
+ def _wrapped (* args : Any , ** kwargs : Any ) -> Any :
130
158
params = (kwargs .pop ("params" , None ) or {}).copy ()
131
159
headers = {
132
160
k .lower (): v
@@ -165,7 +193,9 @@ def _wrapped(*args, **kwargs):
165
193
return _wrapper
166
194
167
195
168
- def _bulk_body (serializer , body ):
196
+ def _bulk_body (
197
+ serializer : Serializer , body : Union [str , bytes , Collection [Any ]]
198
+ ) -> Union [str , bytes ]:
169
199
# if not passed in a string, serialize items and join by newline
170
200
if not isinstance (body , string_types ):
171
201
body = "\n " .join (map (serializer .dumps , body ))
@@ -174,13 +204,15 @@ def _bulk_body(serializer, body):
174
204
if isinstance (body , bytes ):
175
205
if not body .endswith (b"\n " ):
176
206
body += b"\n "
177
- elif isinstance (body , string_types ) and not body .endswith ("\n " ):
207
+ elif isinstance (body , str ) and not body .endswith ("\n " ):
178
208
body += "\n "
179
209
180
210
return body
181
211
182
212
183
- def _base64_auth_header (auth_value ):
213
+ def _base64_auth_header (
214
+ auth_value : Union [List [str ], Tuple [str , ...], str , bytes ]
215
+ ) -> str :
184
216
"""Takes either a 2-tuple or a base64-encoded string
185
217
and returns a base64-encoded string to be used
186
218
as an HTTP authorization header.
@@ -191,17 +223,19 @@ def _base64_auth_header(auth_value):
191
223
192
224
193
225
class NamespacedClient :
194
- def __init__ (self , client ):
226
+ client : "Elasticsearch"
227
+
228
+ def __init__ (self , client : "Elasticsearch" ) -> None :
195
229
self .client = client
196
230
197
231
@property
198
- def transport (self ):
232
+ def transport (self ) -> Transport :
199
233
return self .client .transport
200
234
201
235
202
236
class AddonClient (NamespacedClient ):
203
237
@classmethod
204
- def infect_client (cls , client ) :
238
+ def infect_client (cls , client : "Elasticsearch" ) -> "Elasticsearch" :
205
239
addon = cls (weakref .proxy (client ))
206
- setattr (client , cls .namespace , addon )
240
+ setattr (client , cls .namespace , addon ) # type: ignore
207
241
return client
0 commit comments