[7.x] Add typing metadata and manual types to package

MANIFEST.in

@@ -7,6 +7,7 @@ include README.rst
 include README
 include tox.ini
 include setup.py
+recursive-include elasticsearch* py.typed
 recursive-include docs *
 recursive-exclude docs/examples *
 

docs/sphinx/async.rst

@@ -230,6 +230,12 @@ AsyncTransport
  .. autoclass:: AsyncTransport
    :members:
 
+AsyncConnection
+~~~~~~~~~~~~~~~~~
+
+ .. autoclass:: AsyncConnection
+   :members:
+
 AIOHttpConnection
 ~~~~~~~~~~~~~~~~~
 

elasticsearch/__init__.py

@@ -84,10 +84,15 @@ try:
     if sys.version_info < (3, 6):
         raise ImportError
 
-    from ._async.http_aiohttp import AIOHttpConnection
+    from ._async.http_aiohttp import AIOHttpConnection, AsyncConnection
     from ._async.transport import AsyncTransport
     from ._async.client import AsyncElasticsearch
 
-    __all__ += ["AIOHttpConnection", "AsyncTransport", "AsyncElasticsearch"]
+    __all__ += [
+        "AIOHttpConnection",
+        "AsyncConnection",
+        "AsyncTransport",
+        "AsyncElasticsearch",
+    ]
 except (ImportError, SyntaxError):
     pass

elasticsearch/__init__.pyi

@@ -0,0 +1,62 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import sys
+from typing import Tuple
+
+from .client import Elasticsearch as Elasticsearch
+from .transport import Transport as Transport
+from .connection_pool import (
+    ConnectionPool as ConnectionPool,
+    ConnectionSelector as ConnectionSelector,
+    RoundRobinSelector as RoundRobinSelector,
+)
+from .serializer import JSONSerializer as JSONSerializer
+from .connection import (
+    Connection as Connection,
+    RequestsHttpConnection as RequestsHttpConnection,
+    Urllib3HttpConnection as Urllib3HttpConnection,
+)
+from .exceptions import (
+    ImproperlyConfigured as ImproperlyConfigured,
+    ElasticsearchException as ElasticsearchException,
+    SerializationError as SerializationError,
+    TransportError as TransportError,
+    NotFoundError as NotFoundError,
+    ConflictError as ConflictError,
+    RequestError as RequestError,
+    ConnectionError as ConnectionError,
+    SSLError as SSLError,
+    ConnectionTimeout as ConnectionTimeout,
+    AuthenticationException as AuthenticationException,
+    AuthorizationException as AuthorizationException,
+    ElasticsearchDeprecationWarning as ElasticsearchDeprecationWarning,
+)
+
+try:
+    if sys.version_info < (3, 6):
+        raise ImportError
+
+    from ._async.http_aiohttp import AIOHttpConnection as AIOHttpConnection
+    from ._async.transport import AsyncTransport as AsyncTransport
+    from ._async.client import AsyncElasticsearch as AsyncElasticsearch
+except (ImportError, SyntaxError):
+    pass
+
+VERSION: Tuple[int, int, int]
+__version__: Tuple[int, int, int]
+__versionstr__: str

elasticsearch/_async/_extra_imports.py

@@ -0,0 +1,36 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+# type: ignore
+
+# This file exists for the sole reason of making mypy not
+# complain about type issues to do with 'aiohttp' and 'yarl'.
+# We're in a catch-22 situation:
+# - If we use 'type: ignore' on 'import aiohttp' and it's not installed
+#   mypy will complain that the annotation is unnecessary.
+# - If we don't use 'type: ignore' on 'import aiohttp' and it
+#   it's not installed mypy will complain that it can't find
+#   type hints for aiohttp.
+# So to make mypy happy we move all our 'extra' imports here
+# and add a global 'type: ignore' which mypy never complains
+# about being unnecessary.
+
+import aiohttp
+import aiohttp.client_exceptions as aiohttp_exceptions
+import yarl
+
+__all__ = ["aiohttp", "aiohttp_exceptions", "yarl"]

elasticsearch/_async/compat.pyi

@@ -0,0 +1,20 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import asyncio
+
+def get_running_loop() -> asyncio.AbstractEventLoop: ...

elasticsearch/_async/helpers.pyi

@@ -0,0 +1,101 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import (
+    AsyncIterator,
+    AsyncGenerator,
+    Optional,
+    Union,
+    Any,
+    Mapping,
+    Tuple,
+    TypeVar,
+    Iterable,
+    AsyncIterable,
+    List,
+    Collection,
+    Callable,
+)
+import logging
+from .client import AsyncElasticsearch
+from ..serializer import Serializer
+
+logger: logging.Logger
+
+T = TypeVar("T")
+
+def _chunk_actions(
+    actions: Any, chunk_size: int, max_chunk_bytes: int, serializer: Serializer
+) -> AsyncGenerator[Any, None]: ...
+def _process_bulk_chunk(
+    client: AsyncElasticsearch,
+    bulk_actions: Any,
+    bulk_data: Any,
+    raise_on_exception: bool = ...,
+    raise_on_error: bool = ...,
+    *args: Any,
+    **kwargs: Any
+) -> AsyncGenerator[Tuple[bool, Any], None]: ...
+def aiter(x: Union[Iterable[T], AsyncIterable[T]]) -> AsyncGenerator[T, None]: ...
+def azip(
+    *iterables: Union[Iterable[T], AsyncIterable[T]]
+) -> AsyncGenerator[Tuple[T, ...], None]: ...
+def async_streaming_bulk(
+    client: AsyncElasticsearch,
+    actions: Union[Iterable[Any], AsyncIterable[Any]],
+    chunk_size: int = ...,
+    max_chunk_bytes: int = ...,
+    raise_on_error: bool = ...,
+    expand_action_callback: Callable[[Any], Tuple[str, Any]] = ...,
+    raise_on_exception: bool = ...,
+    max_retries: int = ...,
+    initial_backoff: Union[float, int] = ...,
+    max_backoff: Union[float, int] = ...,
+    yield_ok: bool = ...,
+    *args: Any,
+    **kwargs: Any
+) -> AsyncGenerator[Tuple[bool, Any], None]: ...
+async def async_bulk(
+    client: AsyncElasticsearch,
+    actions: Union[Iterable[Any], AsyncIterable[Any]],
+    stats_only: bool = ...,
+    *args: Any,
+    **kwargs: Any
+) -> Tuple[int, Union[int, List[Any]]]: ...
+def async_scan(
+    client: AsyncElasticsearch,
+    query: Optional[Any] = ...,
+    scroll: str = ...,
+    raise_on_error: bool = ...,
+    preserve_order: bool = ...,
+    size: int = ...,
+    request_timeout: Optional[Union[float, int]] = ...,
+    clear_scroll: bool = ...,
+    scroll_kwargs: Optional[Mapping[str, Any]] = ...,
+    **kwargs: Any
+) -> AsyncGenerator[int, None]: ...
+async def async_reindex(
+    client: AsyncElasticsearch,
+    source_index: Union[str, Collection[str]],
+    target_index: str,
+    query: Any = ...,
+    target_client: Optional[AsyncElasticsearch] = ...,
+    chunk_size: int = ...,
+    scroll: str = ...,
+    scan_kwargs: Optional[Mapping[str, Any]] = ...,
+    bulk_kwargs: Optional[Mapping[str, Any]] = ...,
+) -> Tuple[int, Union[int, List[Any]]]: ...

elasticsearch/_async/http_aiohttp.py

@@ -18,13 +18,9 @@
 import asyncio
 import ssl
 import os
-import urllib3
+import urllib3  # type: ignore
 import warnings
-
+from ._extra_imports import aiohttp_exceptions, aiohttp, yarl
-import aiohttp
-import yarl
-from aiohttp.client_exceptions import ServerFingerprintMismatch, ServerTimeoutError
-
 from .compat import get_running_loop
 from ..connection import Connection
 from ..compat import urlencode
@@ -52,7 +48,26 @@ except ImportError:
     pass
 
 
-class AIOHttpConnection(Connection):
+class AsyncConnection(Connection):
+    """Base class for Async HTTP connection implementations"""
+
+    async def perform_request(
+        self,
+        method,
+        url,
+        params=None,
+        body=None,
+        timeout=None,
+        ignore=(),
+        headers=None,
+    ):
+        raise NotImplementedError()
+
+    async def close(self):
+        raise NotImplementedError()
+
+
+class AIOHttpConnection(AsyncConnection):
     def __init__(
         self,
         host="localhost",
@@ -199,6 +214,7 @@ class AIOHttpConnection(Connection):
     ):
         if self.session is None:
             await self._create_aiohttp_session()
+        assert self.session is not None
 
         orig_body = body
         url_path = self.url_prefix + url
@@ -260,11 +276,18 @@ class AIOHttpConnection(Connection):
 
         except Exception as e:
             self.log_request_fail(
-                method, url, url_path, orig_body, self.loop.time() - start, exception=e
+                method,
+                str(url),
+                url_path,
+                orig_body,
+                self.loop.time() - start,
+                exception=e,
             )
-            if isinstance(e, ServerFingerprintMismatch):
+            if isinstance(e, aiohttp_exceptions.ServerFingerprintMismatch):
                 raise SSLError("N/A", str(e), e)
-            if isinstance(e, (asyncio.TimeoutError, ServerTimeoutError)):
+            if isinstance(
+                e, (asyncio.TimeoutError, aiohttp_exceptions.ServerTimeoutError)
+            ):
                 raise ConnectionTimeout("TIMEOUT", str(e), e)
             raise ConnectionError("N/A", str(e), e)
 
@@ -276,7 +299,7 @@ class AIOHttpConnection(Connection):
         if not (200 <= response.status < 300) and response.status not in ignore:
             self.log_request_fail(
                 method,
-                url,
+                str(url),
                 url_path,
                 orig_body,
                 duration,
@@ -286,7 +309,7 @@ class AIOHttpConnection(Connection):
             self._raise_error(response.status, raw_data)
 
         self.log_request_success(
-            method, url, url_path, orig_body, response.status, raw_data, duration
+            method, str(url), url_path, orig_body, response.status, raw_data, duration
         )
 
         return response.status, response.headers, raw_data
@@ -312,9 +335,7 @@ class AIOHttpConnection(Connection):
             cookie_jar=aiohttp.DummyCookieJar(),
             response_class=ESClientResponse,
             connector=aiohttp.TCPConnector(
-                limit=self._limit,
+                limit=self._limit, use_dns_cache=True, ssl=self._ssl_context
-                use_dns_cache=True,
-                ssl=self._ssl_context,
             ),
         )
 

elasticsearch/_async/http_aiohttp.pyi

@@ -0,0 +1,60 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from ._extra_imports import aiohttp  # type: ignore
+from typing import Optional, Mapping, Collection, Union, Any, Tuple
+from ..connection import Connection
+
+class AsyncConnection(Connection):
+    async def perform_request(  # type: ignore
+        self,
+        method: str,
+        url: str,
+        params: Optional[Mapping[str, Any]] = ...,
+        body: Optional[bytes] = ...,
+        timeout: Optional[Union[int, float]] = ...,
+        ignore: Collection[int] = ...,
+        headers: Optional[Mapping[str, str]] = ...,
+    ) -> Tuple[int, Mapping[str, str], str]: ...
+    async def close(self) -> None: ...
+
+class AIOHttpConnection(AsyncConnection):
+    session: Optional[aiohttp.ClientSession]
+    ssl_assert_fingerprint: Optional[str]
+    def __init__(
+        self,
+        host: str = ...,
+        port: Optional[int] = ...,
+        http_auth: Optional[Any] = ...,
+        use_ssl: bool = ...,
+        verify_certs: bool = ...,
+        ssl_show_warn: bool = ...,
+        ca_certs: Optional[Any] = ...,
+        client_cert: Optional[Any] = ...,
+        client_key: Optional[Any] = ...,
+        ssl_version: Optional[Any] = ...,
+        ssl_assert_fingerprint: Optional[Any] = ...,
+        maxsize: int = ...,
+        headers: Optional[Mapping[str, str]] = ...,
+        ssl_context: Optional[Any] = ...,
+        http_compress: Optional[bool] = ...,
+        cloud_id: Optional[str] = ...,
+        api_key: Optional[Any] = ...,
+        opaque_id: Optional[str] = ...,
+        loop: Any = ...,
+        **kwargs: Any,
+    ) -> None: ...

elasticsearch/_async/transport.pyi

@@ -0,0 +1,91 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import (
+    Callable,
+    Optional,
+    Union,
+    Collection,
+    Type,
+    Mapping,
+    Any,
+    Dict,
+    List,
+)
+
+from ..connection import Connection
+from ..connection_pool import ConnectionPool
+from ..serializer import Serializer, Deserializer
+
+class AsyncTransport(object):
+    DEFAULT_CONNECTION_CLASS: Type[Connection]
+    connection_pool: ConnectionPool
+    deserializer: Deserializer
+
+    max_retries: int
+    retry_on_timeout: bool
+    retry_on_status: Collection[int]
+    send_get_body_as: str
+    serializer: Serializer
+    connection_pool_class: Type[ConnectionPool]
+    connection_class: Type[Connection]
+    kwargs: Any
+    hosts: Optional[List[Dict[str, Any]]]
+    seed_connections: List[Connection]
+    sniffer_timeout: Optional[float]
+    sniff_on_start: bool
+    sniff_on_connection_fail: bool
+    last_sniff: float
+    sniff_timeout: Optional[float]
+    host_info_callback: Callable[
+        [Dict[str, Any], Optional[Dict[str, Any]]], Dict[str, Any]
+    ]
+    def __init__(
+        self,
+        hosts: Any,
+        connection_class: Optional[Type[Any]] = ...,
+        connection_pool_class: Type[ConnectionPool] = ...,
+        host_info_callback: Callable[
+            [Dict[str, Any], Dict[str, Any]], Optional[Dict[str, Any]]
+        ] = ...,
+        sniff_on_start: bool = ...,
+        sniffer_timeout: Optional[Union[float, int]] = ...,
+        sniff_timeout: Union[float, int] = ...,
+        sniff_on_connection_fail: bool = ...,
+        serializer: Serializer = ...,
+        serializers: Optional[Mapping[str, Serializer]] = ...,
+        default_mimetype: str = ...,
+        max_retries: int = ...,
+        retry_on_status: Collection[int] = ...,
+        retry_on_timeout: bool = ...,
+        send_get_body_as: str = ...,
+        **kwargs: Any
+    ) -> None: ...
+    def add_connection(self, host: Any) -> None: ...
+    def set_connections(self, hosts: Collection[Any]) -> None: ...
+    def get_connection(self) -> Connection: ...
+    def sniff_hosts(self, initial: bool = ...) -> None: ...
+    def mark_dead(self, connection: Connection) -> None: ...
+    async def perform_request(
+        self,
+        method: str,
+        url: str,
+        headers: Optional[Mapping[str, str]] = ...,
+        params: Optional[Mapping[str, Any]] = ...,
+        body: Optional[Any] = ...,
+    ) -> Union[bool, str]: ...
+    async def close(self) -> None: ...

elasticsearch/compat.pyi

@@ -0,0 +1,44 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import sys
+from typing import Tuple
+
+PY2: bool
+string_types: Tuple[type, ...]
+
+if sys.version_info[0] == 2:
+    from urllib import (
+        quote_plus as quote_plus,
+        quote as quote,
+        urlencode as urlencode,
+        unquote as unquote,
+    )
+    from urlparse import urlparse as urlparse
+    from itertools import imap as map
+    from Queue import Queue as Queue
+else:
+    from urllib.parse import (
+        quote as quote,
+        quote_plus as quote_plus,
+        urlencode as urlencode,
+        urlparse as urlparse,
+        unquote as unquote,
+    )
+
+    map = map
+    from queue import Queue as Queue

elasticsearch/connection/__init__.pyi

@@ -0,0 +1,23 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from .base import Connection as Connection
+from .http_requests import RequestsHttpConnection as RequestsHttpConnection
+from .http_urllib3 import (
+    Urllib3HttpConnection as Urllib3HttpConnection,
+    create_ssl_context as create_ssl_context,
+)

elasticsearch/connection/base.py

@@ -223,6 +223,18 @@ class Connection(object):
                 self._pretty_json(response).replace("\n", "\n#") if response else "",
             )
 
+    def perform_request(
+        self,
+        method,
+        url,
+        params=None,
+        body=None,
+        timeout=None,
+        ignore=(),
+        headers=None,
+    ):
+        raise NotImplementedError()
+
     def log_request_success(
         self, method, full_url, path, body, status_code, response, duration
     ):

elasticsearch/connection/base.pyi

@@ -0,0 +1,109 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import logging
+
+from typing import (
+    Union,
+    Optional,
+    Mapping,
+    Tuple,
+    List,
+    NoReturn,
+    Dict,
+    Sequence,
+    Any,
+    AnyStr,
+    Collection,
+)
+
+logger: logging.Logger
+tracer: logging.Logger
+
+class Connection(object):
+    headers: Dict[str, str]
+    use_ssl: bool
+    http_compress: bool
+    scheme: str
+    hostname: str
+    port: Optional[int]
+    host: str
+    url_prefix: str
+    timeout: Optional[Union[float, int]]
+    def __init__(
+        self,
+        host: str = ...,
+        port: Optional[int] = ...,
+        use_ssl: bool = ...,
+        url_prefix: str = ...,
+        timeout: Optional[Union[float, int]] = ...,
+        headers: Optional[Mapping[str, str]] = ...,
+        http_compress: Optional[bool] = ...,
+        cloud_id: Optional[str] = ...,
+        api_key: Optional[Union[Tuple[str, str], List[str], str]] = ...,
+        opaque_id: Optional[str] = ...,
+        **kwargs: Any
+    ) -> None: ...
+    def __repr__(self) -> str: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def _gzip_compress(self, body: bytes) -> bytes: ...
+    def _raise_warnings(self, warning_headers: Sequence[str]) -> None: ...
+    def _pretty_json(self, data: Any) -> str: ...
+    def _log_trace(
+        self,
+        method: Any,
+        path: Any,
+        body: Any,
+        status_code: Any,
+        response: Any,
+        duration: Any,
+    ) -> None: ...
+    def perform_request(
+        self,
+        method: str,
+        url: str,
+        params: Optional[Mapping[str, Any]] = ...,
+        body: Optional[bytes] = ...,
+        timeout: Optional[Union[int, float]] = ...,
+        ignore: Collection[int] = ...,
+        headers: Optional[Mapping[str, str]] = ...,
+    ) -> Tuple[int, Mapping[str, str], str]: ...
+    def log_request_success(
+        self,
+        method: str,
+        full_url: str,
+        path: str,
+        body: str,
+        status_code: int,
+        response: bytes,
+        duration: float,
+    ) -> None: ...
+    def log_request_fail(
+        self,
+        method: str,
+        full_url: str,
+        path: str,
+        body: str,
+        duration: float,
+        status_code: Optional[int] = ...,
+        response: Optional[bytes] = ...,
+        exception: Optional[Exception] = ...,
+    ) -> None: ...
+    def _raise_error(self, status_code: int, raw_data: str) -> NoReturn: ...
+    def _get_default_user_agent(self) -> str: ...
+    def _get_api_key_header_val(self, api_key: Any) -> str: ...

elasticsearch/connection/http_requests.pyi

@@ -0,0 +1,41 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import Optional, Any, Mapping
+import requests
+from .base import Connection
+
+class RequestsHttpConnection(Connection):
+    session: requests.Session
+    def __init__(
+        self,
+        host: str = ...,
+        port: Optional[int] = ...,
+        http_auth: Optional[Any] = ...,
+        use_ssl: bool = ...,
+        verify_certs: bool = ...,
+        ssl_show_warn: bool = ...,
+        ca_certs: Optional[Any] = ...,
+        client_cert: Optional[Any] = ...,
+        client_key: Optional[Any] = ...,
+        headers: Optional[Mapping[str, str]] = ...,
+        http_compress: Optional[bool] = ...,
+        cloud_id: Optional[str] = ...,
+        api_key: Optional[Any] = ...,
+        opaque_id: Optional[str] = ...,
+        **kwargs: Any
+    ) -> None: ...

elasticsearch/connection/http_urllib3.py

@@ -17,9 +17,9 @@
 
 import time
 import ssl
-import urllib3
+import urllib3  # type: ignore
-from urllib3.exceptions import ReadTimeoutError, SSLError as UrllibSSLError
+from urllib3.exceptions import ReadTimeoutError, SSLError as UrllibSSLError  # type: ignore
-from urllib3.util.retry import Retry
+from urllib3.util.retry import Retry  # type: ignore
 import warnings
 
 from .base import Connection

elasticsearch/connection/http_urllib3.pyi

@@ -0,0 +1,55 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import ssl
+from typing import Optional, Mapping, Any, Union
+import urllib3  # type: ignore
+from .base import Connection
+
+def create_ssl_context(
+    cafile: Any = ...,
+    capath: Any = ...,
+    cadata: Any = ...,
+) -> ssl.SSLContext: ...
+
+class Urllib3HttpConnection(Connection):
+    pool: urllib3.HTTPConnectionPool
+    def __init__(
+        self,
+        host: str = ...,
+        port: Optional[int] = ...,
+        url_prefix: str = ...,
+        timeout: Optional[Union[float, int]] = ...,
+        http_auth: Any = ...,
+        use_ssl: bool = ...,
+        verify_certs: bool = ...,
+        ssl_show_warn: bool = ...,
+        ca_certs: Optional[Any] = ...,
+        client_cert: Optional[Any] = ...,
+        client_key: Optional[Any] = ...,
+        ssl_version: Optional[Any] = ...,
+        ssl_assert_hostname: Optional[Any] = ...,
+        ssl_assert_fingerprint: Optional[Any] = ...,
+        maxsize: int = ...,
+        headers: Optional[Mapping[str, str]] = ...,
+        ssl_context: Optional[Any] = ...,
+        http_compress: Optional[bool] = ...,
+        cloud_id: Optional[str] = ...,
+        api_key: Optional[Any] = ...,
+        opaque_id: Optional[str] = ...,
+        **kwargs: Any
+    ) -> None: ...

elasticsearch/connection/pooling.py

@@ -15,11 +15,12 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+from .base import Connection
+
 try:
     import queue
 except ImportError:
-    import Queue as queue
+    import Queue as queue  # type: ignore
-from .base import Connection
 
 
 class PoolingConnection(Connection):
@@ -34,6 +35,9 @@ class PoolingConnection(Connection):
         self._free_connections = queue.Queue()
         super(PoolingConnection, self).__init__(*args, **kwargs)
 
+    def _make_connection(self):
+        raise NotImplementedError
+
     def _get_connection(self):
         try:
             return self._free_connections.get_nowait()

elasticsearch/connection/pooling.pyi

@@ -0,0 +1,24 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from .base import Connection
+
+class PoolingConnection(Connection):
+    def _make_connection(self) -> Connection: ...
+    def _get_connection(self) -> Connection: ...
+    def _release_connection(self, con: Connection) -> None: ...
+    def close(self) -> None: ...

elasticsearch/connection_pool.pyi

@@ -0,0 +1,75 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import logging
+from typing import Sequence, Optional, Type, Any, Union, List, Tuple, Dict
+from .connection import Connection
+
+try:
+    from Queue import PriorityQueue  # type: ignore
+except ImportError:
+    from queue import PriorityQueue
+
+logger: logging.Logger
+
+class ConnectionSelector(object):
+    connection_opts: Sequence[Tuple[Connection, Any]]
+    def __init__(self, opts: Sequence[Tuple[Connection, Any]]) -> None: ...
+    def select(self, connections: Sequence[Connection]) -> Connection: ...
+
+class RandomSelector(ConnectionSelector): ...
+class RoundRobinSelector(ConnectionSelector): ...
+
+class ConnectionPool(object):
+    connections_opts: Sequence[Tuple[Connection, Any]]
+    connections: Sequence[Connection]
+    orig_connections: Tuple[Connection, ...]
+    dead: PriorityQueue
+    dead_count: Dict[Connection, int]
+    dead_timeout: Union[float, int]
+    timeout_cutoff: Union[float, int]
+    selector: ConnectionSelector
+    def __init__(
+        self,
+        connections: Sequence[Tuple[Connection, Any]],
+        dead_timeout: int = ...,
+        timeout_cutoff: int = ...,
+        selector_class: Type[ConnectionSelector] = ...,
+        randomize_hosts: bool = ...,
+        **kwargs: Any
+    ) -> None: ...
+    def mark_dead(self, connection: Connection, now: Optional[float] = ...) -> None: ...
+    def mark_live(self, connection: Connection) -> None: ...
+    def resurrect(self, force: bool = ...) -> Optional[Connection]: ...
+    def get_connection(self) -> Connection: ...
+    def close(self) -> None: ...
+    def __repr__(self) -> str: ...
+
+class DummyConnectionPool(ConnectionPool):
+    def __init__(
+        self, connections: Sequence[Tuple[Connection, Any]], **kwargs: Any
+    ) -> None: ...
+    def get_connection(self) -> Connection: ...
+    def close(self) -> None: ...
+    def _noop(self, *args: Any, **kwargs: Any) -> Any: ...
+    mark_dead = mark_live = resurrect = _noop
+
+class EmptyConnectionPool(ConnectionPool):
+    def __init__(self, *_: Any, **__: Any) -> None: ...
+    def get_connection(self) -> Connection: ...
+    def _noop(self, *args: Any, **kwargs: Any) -> Any: ...
+    close = mark_dead = mark_live = resurrect = _noop

elasticsearch/exceptions.pyi

@@ -0,0 +1,48 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import Dict, Union, Any
+
+class ImproperlyConfigured(Exception): ...
+class ElasticsearchException(Exception): ...
+class SerializationError(ElasticsearchException): ...
+
+class TransportError(ElasticsearchException):
+    @property
+    def status_code(self) -> Union[str, int]: ...
+    @property
+    def error(self) -> str: ...
+    @property
+    def info(self) -> Union[Dict[str, Any], Exception]: ...
+    def __str__(self) -> str: ...
+
+class ConnectionError(TransportError):
+    def __str__(self) -> str: ...
+
+class SSLError(ConnectionError): ...
+
+class ConnectionTimeout(ConnectionError):
+    def __str__(self) -> str: ...
+
+class NotFoundError(TransportError): ...
+class ConflictError(TransportError): ...
+class RequestError(TransportError): ...
+class AuthenticationException(TransportError): ...
+class AuthorizationException(TransportError): ...
+class ElasticsearchDeprecationWarning(Warning): ...
+
+HTTP_EXCEPTIONS: Dict[int, ElasticsearchException]

elasticsearch/helpers/__init__.pyi

@@ -0,0 +1,43 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import sys
+from .errors import BulkIndexError as BulkIndexError, ScanError as ScanError
+from .actions import (
+    expand_action as expand_action,
+    streaming_bulk as streaming_bulk,
+    bulk as bulk,
+    parallel_bulk as parallel_bulk,
+    scan as scan,
+    reindex as reindex,
+    _chunk_actions as _chunk_actions,
+    _process_bulk_chunk as _process_bulk_chunk,
+)
+
+try:
+    # Asyncio only supported on Python 3.6+
+    if sys.version_info < (3, 6):
+        raise ImportError
+
+    from .._async.helpers import (
+        async_scan as async_scan,
+        async_bulk as async_bulk,
+        async_reindex as async_reindex,
+        async_streaming_bulk as async_streaming_bulk,
+    )
+except (ImportError, SyntaxError):
+    pass

elasticsearch/helpers/actions.py

@@ -436,7 +436,7 @@ def parallel_bulk(
 
     class BlockingPool(ThreadPool):
         def _setup_queues(self):
-            super(BlockingPool, self)._setup_queues()
+            super(BlockingPool, self)._setup_queues()  # type: ignore
             # The queue must be at least the size of the number of threads to
             # prevent hanging when inserting sentinel values during teardown.
             self._inqueue = Queue(max(queue_size, thread_count))

elasticsearch/helpers/actions.pyi

@@ -0,0 +1,105 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import (
+    Generator,
+    Optional,
+    Union,
+    Any,
+    Mapping,
+    Tuple,
+    Iterable,
+    AsyncIterable,
+    List,
+    Collection,
+    Callable,
+)
+import logging
+from ..client import Elasticsearch
+from ..serializer import Serializer
+
+logger: logging.Logger
+
+def expand_action(data: Any) -> Tuple[str, Any]: ...
+def _chunk_actions(
+    actions: Any, chunk_size: int, max_chunk_bytes: int, serializer: Serializer
+) -> Generator[Any, None, None]: ...
+def _process_bulk_chunk(
+    client: Elasticsearch,
+    bulk_actions: Any,
+    bulk_data: Any,
+    raise_on_exception: bool = ...,
+    raise_on_error: bool = ...,
+    *args: Any,
+    **kwargs: Any
+) -> Generator[Tuple[bool, Any], None, None]: ...
+def streaming_bulk(
+    client: Elasticsearch,
+    actions: Union[Iterable[Any], AsyncIterable[Any]],
+    chunk_size: int = ...,
+    max_chunk_bytes: int = ...,
+    raise_on_error: bool = ...,
+    expand_action_callback: Callable[[Any], Tuple[str, Any]] = ...,
+    raise_on_exception: bool = ...,
+    max_retries: int = ...,
+    initial_backoff: Union[float, int] = ...,
+    max_backoff: Union[float, int] = ...,
+    yield_ok: bool = ...,
+    *args: Any,
+    **kwargs: Any
+) -> Generator[Tuple[bool, Any], None, None]: ...
+def bulk(
+    client: Elasticsearch,
+    actions: Iterable[Any],
+    stats_only: bool = ...,
+    *args: Any,
+    **kwargs: Any
+) -> Tuple[int, Union[int, List[Any]]]: ...
+def parallel_bulk(
+    client: Elasticsearch,
+    actions: Iterable[Any],
+    thread_count: int = ...,
+    chunk_size: int = ...,
+    max_chunk_bytes: int = ...,
+    queue_size: int = ...,
+    expand_action_callback: Callable[[Any], Tuple[str, Any]] = ...,
+    *args: Any,
+    **kwargs: Any
+) -> Generator[Tuple[bool, Any], None, None]: ...
+def scan(
+    client: Elasticsearch,
+    query: Optional[Any] = ...,
+    scroll: str = ...,
+    raise_on_error: bool = ...,
+    preserve_order: bool = ...,
+    size: int = ...,
+    request_timeout: Optional[Union[float, int]] = ...,
+    clear_scroll: bool = ...,
+    scroll_kwargs: Optional[Mapping[str, Any]] = ...,
+    **kwargs: Any
+) -> Generator[Any, None, None]: ...
+def reindex(
+    client: Elasticsearch,
+    source_index: Union[str, Collection[str]],
+    target_index: str,
+    query: Any = ...,
+    target_client: Optional[Elasticsearch] = ...,
+    chunk_size: int = ...,
+    scroll: str = ...,
+    scan_kwargs: Optional[Mapping[str, Any]] = ...,
+    bulk_kwargs: Optional[Mapping[str, Any]] = ...,
+) -> Tuple[int, Union[int, List[Any]]]: ...

elasticsearch/helpers/errors.py

@@ -27,5 +27,5 @@ class BulkIndexError(ElasticsearchException):
 
 class ScanError(ElasticsearchException):
     def __init__(self, scroll_id, *args, **kwargs):
-        super(ScanError, self).__init__(*args, **kwargs)
+        super(ScanError, self).__init__(*args, **kwargs)  # type: ignore
         self.scroll_id = scroll_id

elasticsearch/helpers/errors.pyi

@@ -0,0 +1,27 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import List, Any
+from ..exceptions import ElasticsearchException
+
+class BulkIndexError(ElasticsearchException):
+    @property
+    def errors(self) -> List[Any]: ...
+
+class ScanError(ElasticsearchException):
+    scroll_id: str
+    def __init__(self, scroll_id: str, *args: Any, **kwargs: Any) -> None: ...

elasticsearch/helpers/test.py

@@ -15,6 +15,8 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+# type: ignore
+
 import time
 import os
 from unittest import TestCase, SkipTest

elasticsearch/helpers/test.pyi

@@ -0,0 +1,31 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import Any, Tuple
+from unittest import TestCase
+from ..client import Elasticsearch
+
+def get_test_client(nowait: bool = ..., **kwargs: Any) -> Elasticsearch: ...
+def _get_version(version_string: str) -> Tuple[int, ...]: ...
+
+class ElasticsearchTestCase(TestCase):
+    @staticmethod
+    def _get_client() -> Elasticsearch: ...
+    @classmethod
+    def setup_class(cls) -> None: ...
+    def teardown_method(self, _: Any) -> None: ...
+    def es_version(self) -> Tuple[int, ...]: ...

elasticsearch/serializer.py

@@ -63,7 +63,17 @@ except ImportError:
     pd = None
 
 
-class TextSerializer(object):
+class Serializer(object):
+    mimetype = ""
+
+    def loads(self, s):
+        raise NotImplementedError
+
+    def dumps(self, data):
+        raise NotImplementedError()
+
+
+class TextSerializer(Serializer):
     mimetype = "text/plain"
 
     def loads(self, s):
@@ -76,7 +86,7 @@ class TextSerializer(object):
         raise SerializationError("Cannot serialize %r into text." % data)
 
 
-class JSONSerializer(object):
+class JSONSerializer(Serializer):
     mimetype = "application/json"
 
     def default(self, data):

elasticsearch/serializer.pyi

@@ -0,0 +1,44 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import Optional, Any, Dict
+
+class Serializer(object):
+    mimetype: str
+    def loads(self, s: str) -> Any: ...
+    def dumps(self, data: Any) -> str: ...
+
+class TextSerializer(Serializer):
+    mimetype: str
+    def loads(self, s: str) -> Any: ...
+    def dumps(self, data: Any) -> str: ...
+
+class JSONSerializer(Serializer):
+    mimetype: str
+    def default(self, data: Any) -> Any: ...
+    def loads(self, s: str) -> Any: ...
+    def dumps(self, data: Any) -> str: ...
+
+DEFAULT_SERIALIZERS: Dict[str, Serializer]
+
+class Deserializer(object):
+    def __init__(
+        self,
+        serializers: Dict[str, Serializer],
+        default_mimetype: str = ...,
+    ) -> None: ...
+    def loads(self, s: str, mimetype: Optional[str] = ...) -> Any: ...

elasticsearch/transport.pyi

@@ -0,0 +1,95 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import (
+    Callable,
+    Optional,
+    Union,
+    Collection,
+    Type,
+    Mapping,
+    Any,
+    Dict,
+    List,
+)
+
+from .connection import Connection
+from .connection_pool import ConnectionPool
+from .serializer import Serializer, Deserializer
+
+def get_host_info(
+    node_info: Dict[str, Any], host: Optional[Dict[str, Any]]
+) -> Optional[Dict[str, Any]]: ...
+
+class Transport(object):
+    DEFAULT_CONNECTION_CLASS: Type[Connection]
+    connection_pool: ConnectionPool
+    deserializer: Deserializer
+
+    max_retries: int
+    retry_on_timeout: bool
+    retry_on_status: Collection[int]
+    send_get_body_as: str
+    serializer: Serializer
+    connection_pool_class: Type[ConnectionPool]
+    connection_class: Type[Connection]
+    kwargs: Any
+    hosts: Optional[List[Dict[str, Any]]]
+    seed_connections: List[Connection]
+    sniffer_timeout: Optional[float]
+    sniff_on_start: bool
+    sniff_on_connection_fail: bool
+    last_sniff: float
+    sniff_timeout: Optional[float]
+    host_info_callback: Callable[
+        [Dict[str, Any], Optional[Dict[str, Any]]], Dict[str, Any]
+    ]
+    def __init__(
+        self,
+        hosts: Any,
+        connection_class: Optional[Type[Any]] = ...,
+        connection_pool_class: Type[ConnectionPool] = ...,
+        host_info_callback: Callable[
+            [Dict[str, Any], Dict[str, Any]], Optional[Dict[str, Any]]
+        ] = ...,
+        sniff_on_start: bool = ...,
+        sniffer_timeout: Optional[Union[float, int]] = ...,
+        sniff_timeout: Union[float, int] = ...,
+        sniff_on_connection_fail: bool = ...,
+        serializer: Serializer = ...,
+        serializers: Optional[Mapping[str, Serializer]] = ...,
+        default_mimetype: str = ...,
+        max_retries: int = ...,
+        retry_on_status: Collection[int] = ...,
+        retry_on_timeout: bool = ...,
+        send_get_body_as: str = ...,
+        **kwargs: Any
+    ) -> None: ...
+    def add_connection(self, host: Any) -> None: ...
+    def set_connections(self, hosts: Collection[Any]) -> None: ...
+    def get_connection(self) -> Connection: ...
+    def sniff_hosts(self, initial: bool = ...) -> None: ...
+    def mark_dead(self, connection: Connection) -> None: ...
+    def perform_request(
+        self,
+        method: str,
+        url: str,
+        headers: Optional[Mapping[str, str]] = ...,
+        params: Optional[Mapping[str, Any]] = ...,
+        body: Optional[Any] = ...,
+    ) -> Union[str, bool]: ...
+    def close(self) -> None: ...

noxfile.py

@@ -47,16 +47,33 @@ def blacken(session):
 
 @nox.session()
 def lint(session):
-    session.install("flake8", "black")
+    session.install("flake8", "black", "mypy")
 
     session.run("black", "--target-version=py27", "--check", *SOURCE_FILES)
     session.run("flake8", *SOURCE_FILES)
     session.run("python", "utils/license_headers.py", "check", *SOURCE_FILES)
 
+    # Workaround to make '-r' to still work despite uninstalling aiohttp below.
+    session.run("python", "-m", "pip", "install", "aiohttp")
+
+    # Run mypy on the package and then the type examples separately for
+    # the two different mypy use-cases, ourselves and our users.
+    session.run("mypy", "--strict", "elasticsearch/")
+    session.run("mypy", "--strict", "test_elasticsearch/test_types/")
+
+    # Make sure we don't require aiohttp to be installed for users to
+    # receive type hint information from mypy.
+    session.run("python", "-m", "pip", "uninstall", "--yes", "aiohttp", "yarl")
+    session.run("mypy", "--strict", "elasticsearch/")
+    session.run("mypy", "--strict", "test_elasticsearch/test_types/sync_types.py")
+
 
 @nox.session()
 def docs(session):
     session.install(".")
-    session.install("-rdev-requirements.txt", "sphinx-rtd-theme")
+    session.install(
+        "-rdev-requirements.txt", "sphinx-rtd-theme", "sphinx-autodoc-typehints"
+    )
+    session.run("python", "-m", "pip", "install", "sphinx-autodoc-typehints")
 
     session.run("sphinx-build", "docs/sphinx/", "docs/sphinx/_build", "-b", "html")

setup.py

@@ -61,6 +61,9 @@ setup(
         "Issue Tracker": "https://github.com/elastic/elasticsearch-py/issues",
     },
     packages=find_packages(where=".", exclude=("test_elasticsearch*",)),
+    package_data={"elasticsearch": ["py.typed"]},
+    include_package_data=True,
+    zip_safe=False,
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "License :: OSI Approved :: Apache Software License",

test_elasticsearch/test_types/README.md

@@ -0,0 +1,6 @@
+# Type Hints
+
+All of these scripts are used to test the type hinting
+distributed with the `elasticsearch` package.
+These scripts simulate normal usage of the client and are run
+through `mypy --strict` as a part of continuous integration.

test_elasticsearch/test_types/async_types.py

@@ -0,0 +1,109 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import AsyncGenerator, Dict, Any
+from elasticsearch import (
+    AsyncElasticsearch,
+    AsyncTransport,
+    AIOHttpConnection,
+    ConnectionPool,
+)
+from elasticsearch.helpers import (
+    async_scan,
+    async_streaming_bulk,
+    async_reindex,
+    async_bulk,
+)
+
+
+es = AsyncElasticsearch(
+    [{"host": "localhost", "port": 9443}],
+    transport_class=AsyncTransport,
+)
+t = AsyncTransport(
+    [{}],
+    connection_class=AIOHttpConnection,
+    connection_pool_class=ConnectionPool,
+    sniff_on_start=True,
+    sniffer_timeout=0.1,
+    sniff_timeout=1,
+    sniff_on_connection_fail=False,
+    max_retries=1,
+    retry_on_status={100, 400, 503},
+    retry_on_timeout=True,
+    send_get_body_as="source",
+)
+
+
+async def async_gen() -> AsyncGenerator[Dict[Any, Any], None]:
+    yield {}
+
+
+async def async_scan_types() -> None:
+    async for _ in async_scan(
+        es,
+        query={"query": {"match_all": {}}},
+        request_timeout=10,
+        clear_scroll=True,
+        scroll_kwargs={"request_timeout": 10},
+    ):
+        pass
+    async for _ in async_scan(
+        es,
+        raise_on_error=False,
+        preserve_order=False,
+        scroll="10m",
+        size=10,
+        request_timeout=10.0,
+    ):
+        pass
+
+
+async def async_streaming_bulk_types() -> None:
+    async for _ in async_streaming_bulk(es, async_gen()):
+        pass
+    async for _ in async_streaming_bulk(es, async_gen().__aiter__()):
+        pass
+    async for _ in async_streaming_bulk(es, [{}]):
+        pass
+    async for _ in async_streaming_bulk(es, ({},)):
+        pass
+
+
+async def async_bulk_types() -> None:
+    _, _ = await async_bulk(es, async_gen())
+    _, _ = await async_bulk(es, async_gen().__aiter__())
+    _, _ = await async_bulk(es, [{}])
+    _, _ = await async_bulk(es, ({},))
+
+
+async def async_reindex_types() -> None:
+    _, _ = await async_reindex(
+        es, "src-index", "target-index", query={"query": {"match": {"key": "val"}}}
+    )
+    _, _ = await async_reindex(
+        es, source_index="src-index", target_index="target-index", target_client=es
+    )
+    _, _ = await async_reindex(
+        es,
+        "src-index",
+        "target-index",
+        chunk_size=1,
+        scroll="10m",
+        scan_kwargs={"request_timeout": 10},
+        bulk_kwargs={"request_timeout": 10},
+    )

test_elasticsearch/test_types/sync_types.py

@@ -0,0 +1,104 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import Generator, Dict, Any
+from elasticsearch import (
+    Elasticsearch,
+    Transport,
+    RequestsHttpConnection,
+    ConnectionPool,
+)
+from elasticsearch.helpers import scan, streaming_bulk, reindex, bulk
+
+
+es = Elasticsearch(
+    [{"host": "localhost", "port": 9443}],
+    transport_class=Transport,
+)
+t = Transport(
+    [{}],
+    connection_class=RequestsHttpConnection,
+    connection_pool_class=ConnectionPool,
+    sniff_on_start=True,
+    sniffer_timeout=0.1,
+    sniff_timeout=1,
+    sniff_on_connection_fail=False,
+    max_retries=1,
+    retry_on_status={100, 400, 503},
+    retry_on_timeout=True,
+    send_get_body_as="source",
+)
+
+
+def sync_gen() -> Generator[Dict[Any, Any], None, None]:
+    yield {}
+
+
+def scan_types() -> None:
+    for _ in scan(
+        es,
+        query={"query": {"match_all": {}}},
+        request_timeout=10,
+        clear_scroll=True,
+        scroll_kwargs={"request_timeout": 10},
+    ):
+        pass
+    for _ in scan(
+        es,
+        raise_on_error=False,
+        preserve_order=False,
+        scroll="10m",
+        size=10,
+        request_timeout=10.0,
+    ):
+        pass
+
+
+def streaming_bulk_types() -> None:
+    for _ in streaming_bulk(es, sync_gen()):
+        pass
+    for _ in streaming_bulk(es, sync_gen().__iter__()):
+        pass
+    for _ in streaming_bulk(es, [{}]):
+        pass
+    for _ in streaming_bulk(es, ({},)):
+        pass
+
+
+def bulk_types() -> None:
+    _, _ = bulk(es, sync_gen())
+    _, _ = bulk(es, sync_gen().__iter__())
+    _, _ = bulk(es, [{}])
+    _, _ = bulk(es, ({},))
+
+
+def reindex_types() -> None:
+    _, _ = reindex(
+        es, "src-index", "target-index", query={"query": {"match": {"key": "val"}}}
+    )
+    _, _ = reindex(
+        es, source_index="src-index", target_index="target-index", target_client=es
+    )
+    _, _ = reindex(
+        es,
+        "src-index",
+        "target-index",
+        chunk_size=1,
+        scroll="10m",
+        scan_kwargs={"request_timeout": 10},
+        bulk_kwargs={"request_timeout": 10},
+    )

utils/license_headers.py

@@ -21,6 +21,7 @@ error out if 'fix' would have changed the file.
 """
 
 import os
+import re
 import sys
 from typing import List, Iterator
 from itertools import chain
@@ -64,7 +65,7 @@ def find_files_to_fix(sources: List[str]) -> Iterator[str]:
 
 
 def does_file_need_fix(filepath: str) -> bool:
-    if not filepath.endswith(".py"):
+    if not re.search(r"\.pyi?$", filepath):
         return False
     with open(filepath, mode="r") as f:
         first_license_line = None