diff --git a/MANIFEST.in b/MANIFEST.in index fb657b9b..578b5a60 100644 --- a/MANIFEST.in +++ b/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 * diff --git a/docs/sphinx/async.rst b/docs/sphinx/async.rst index fd2654f9..a8853038 100644 --- a/docs/sphinx/async.rst +++ b/docs/sphinx/async.rst @@ -230,6 +230,12 @@ AsyncTransport .. autoclass:: AsyncTransport :members: +AsyncConnection +~~~~~~~~~~~~~~~~~ + + .. autoclass:: AsyncConnection + :members: + AIOHttpConnection ~~~~~~~~~~~~~~~~~ diff --git a/elasticsearch/__init__.py b/elasticsearch/__init__.py index 02892e3b..9eab5b0e 100644 --- a/elasticsearch/__init__.py +++ b/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 diff --git a/elasticsearch/__init__.pyi b/elasticsearch/__init__.pyi new file mode 100644 index 00000000..63cabf85 --- /dev/null +++ b/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 diff --git a/elasticsearch/_async/_extra_imports.py b/elasticsearch/_async/_extra_imports.py new file mode 100644 index 00000000..08c47fc0 --- /dev/null +++ b/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"] diff --git a/elasticsearch/_async/compat.pyi b/elasticsearch/_async/compat.pyi new file mode 100644 index 00000000..fe86b8f7 --- /dev/null +++ b/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: ... diff --git a/elasticsearch/_async/helpers.pyi b/elasticsearch/_async/helpers.pyi new file mode 100644 index 00000000..e8a375c7 --- /dev/null +++ b/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]]]: ... diff --git a/elasticsearch/_async/http_aiohttp.py b/elasticsearch/_async/http_aiohttp.py index 16576a8e..ce76ce14 100644 --- a/elasticsearch/_async/http_aiohttp.py +++ b/elasticsearch/_async/http_aiohttp.py @@ -18,13 +18,9 @@ import asyncio import ssl import os -import urllib3 +import urllib3 # type: ignore import warnings - -import aiohttp -import yarl -from aiohttp.client_exceptions import ServerFingerprintMismatch, ServerTimeoutError - +from ._extra_imports import aiohttp_exceptions, aiohttp, yarl 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, - use_dns_cache=True, - ssl=self._ssl_context, + limit=self._limit, use_dns_cache=True, ssl=self._ssl_context ), ) diff --git a/elasticsearch/_async/http_aiohttp.pyi b/elasticsearch/_async/http_aiohttp.pyi new file mode 100644 index 00000000..971f98ce --- /dev/null +++ b/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: ... diff --git a/elasticsearch/_async/transport.pyi b/elasticsearch/_async/transport.pyi new file mode 100644 index 00000000..13be2247 --- /dev/null +++ b/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: ... diff --git a/elasticsearch/compat.pyi b/elasticsearch/compat.pyi new file mode 100644 index 00000000..00eefc9b --- /dev/null +++ b/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 diff --git a/elasticsearch/connection/__init__.pyi b/elasticsearch/connection/__init__.pyi new file mode 100644 index 00000000..95a73ddd --- /dev/null +++ b/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, +) diff --git a/elasticsearch/connection/base.py b/elasticsearch/connection/base.py index 7228b2cc..edfa72b1 100644 --- a/elasticsearch/connection/base.py +++ b/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 ): diff --git a/elasticsearch/connection/base.pyi b/elasticsearch/connection/base.pyi new file mode 100644 index 00000000..121e22bb --- /dev/null +++ b/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: ... diff --git a/elasticsearch/connection/http_requests.pyi b/elasticsearch/connection/http_requests.pyi new file mode 100644 index 00000000..a9964837 --- /dev/null +++ b/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: ... diff --git a/elasticsearch/connection/http_urllib3.py b/elasticsearch/connection/http_urllib3.py index 9ef5406a..bb2a2a9b 100644 --- a/elasticsearch/connection/http_urllib3.py +++ b/elasticsearch/connection/http_urllib3.py @@ -17,9 +17,9 @@ import time import ssl -import urllib3 -from urllib3.exceptions import ReadTimeoutError, SSLError as UrllibSSLError -from urllib3.util.retry import Retry +import urllib3 # type: ignore +from urllib3.exceptions import ReadTimeoutError, SSLError as UrllibSSLError # type: ignore +from urllib3.util.retry import Retry # type: ignore import warnings from .base import Connection diff --git a/elasticsearch/connection/http_urllib3.pyi b/elasticsearch/connection/http_urllib3.pyi new file mode 100644 index 00000000..47b126b3 --- /dev/null +++ b/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: ... diff --git a/elasticsearch/connection/pooling.py b/elasticsearch/connection/pooling.py index de2ccd35..ea178c51 100644 --- a/elasticsearch/connection/pooling.py +++ b/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 -from .base import Connection + import Queue as queue # type: ignore 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() diff --git a/elasticsearch/connection/pooling.pyi b/elasticsearch/connection/pooling.pyi new file mode 100644 index 00000000..ba6f912c --- /dev/null +++ b/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: ... diff --git a/elasticsearch/connection_pool.pyi b/elasticsearch/connection_pool.pyi new file mode 100644 index 00000000..628e6219 --- /dev/null +++ b/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 diff --git a/elasticsearch/exceptions.pyi b/elasticsearch/exceptions.pyi new file mode 100644 index 00000000..f8c5b02d --- /dev/null +++ b/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] diff --git a/elasticsearch/helpers/__init__.pyi b/elasticsearch/helpers/__init__.pyi new file mode 100644 index 00000000..7a01f782 --- /dev/null +++ b/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 diff --git a/elasticsearch/helpers/actions.py b/elasticsearch/helpers/actions.py index 79693f4b..454696ee 100644 --- a/elasticsearch/helpers/actions.py +++ b/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)) diff --git a/elasticsearch/helpers/actions.pyi b/elasticsearch/helpers/actions.pyi new file mode 100644 index 00000000..20f8a3d6 --- /dev/null +++ b/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]]]: ... diff --git a/elasticsearch/helpers/errors.py b/elasticsearch/helpers/errors.py index e6292cf9..edeffe83 100644 --- a/elasticsearch/helpers/errors.py +++ b/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 diff --git a/elasticsearch/helpers/errors.pyi b/elasticsearch/helpers/errors.pyi new file mode 100644 index 00000000..8dd96fc4 --- /dev/null +++ b/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: ... diff --git a/elasticsearch/helpers/test.py b/elasticsearch/helpers/test.py index bf12041e..662a38fc 100644 --- a/elasticsearch/helpers/test.py +++ b/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 diff --git a/elasticsearch/helpers/test.pyi b/elasticsearch/helpers/test.pyi new file mode 100644 index 00000000..06ff323b --- /dev/null +++ b/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, ...]: ... diff --git a/elasticsearch/py.typed b/elasticsearch/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/elasticsearch/serializer.py b/elasticsearch/serializer.py index c78df0e9..ab19342a 100644 --- a/elasticsearch/serializer.py +++ b/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): diff --git a/elasticsearch/serializer.pyi b/elasticsearch/serializer.pyi new file mode 100644 index 00000000..441ba1e8 --- /dev/null +++ b/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: ... diff --git a/elasticsearch/transport.pyi b/elasticsearch/transport.pyi new file mode 100644 index 00000000..def26c49 --- /dev/null +++ b/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: ... diff --git a/noxfile.py b/noxfile.py index 863926c4..52655103 100644 --- a/noxfile.py +++ b/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") diff --git a/setup.py b/setup.py index 2e8b0b8b..78141d0e 100644 --- a/setup.py +++ b/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", diff --git a/test_elasticsearch/test_types/README.md b/test_elasticsearch/test_types/README.md new file mode 100644 index 00000000..0fc3182e --- /dev/null +++ b/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. diff --git a/test_elasticsearch/test_types/async_types.py b/test_elasticsearch/test_types/async_types.py new file mode 100644 index 00000000..f073d8e9 --- /dev/null +++ b/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}, + ) diff --git a/test_elasticsearch/test_types/sync_types.py b/test_elasticsearch/test_types/sync_types.py new file mode 100644 index 00000000..9767aa00 --- /dev/null +++ b/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}, + ) diff --git a/utils/license_headers.py b/utils/license_headers.py index 217854d1..cf587eba 100644 --- a/utils/license_headers.py +++ b/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