diff --git a/MANIFEST.in b/MANIFEST.in index b0ca5932..549b8295 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -6,7 +6,7 @@ include MANIFEST.in include README.rst include README include setup.py -recursive-include elasticsearch* py.typed +recursive-include elasticsearch* py.typed *.pyi recursive-include docs/sphinx * prune docs/sphinx/_build diff --git a/noxfile.py b/noxfile.py index 8cc5cac8..9df9cf46 100644 --- a/noxfile.py +++ b/noxfile.py @@ -59,7 +59,8 @@ def lint(session): # 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/") + session.run("mypy", "--strict", "test_elasticsearch/test_types/sync_types.py") + session.run("mypy", "--strict", "test_elasticsearch/test_types/async_types.py") # Make sure we don't require aiohttp to be installed for users to # receive type hint information from mypy. diff --git a/setup.py b/setup.py index eccb7661..0bed1436 100644 --- a/setup.py +++ b/setup.py @@ -72,7 +72,7 @@ setup( "Issue Tracker": "https://github.com/elastic/elasticsearch-py/issues", }, packages=packages, - package_data={"elasticsearch": ["py.typed"]}, + package_data={"elasticsearch": ["py.typed", "*.pyi"]}, include_package_data=True, zip_safe=False, classifiers=[ diff --git a/test_elasticsearch/test_types/aliased_types.py b/test_elasticsearch/test_types/aliased_types.py new file mode 100644 index 00000000..45e28306 --- /dev/null +++ b/test_elasticsearch/test_types/aliased_types.py @@ -0,0 +1,195 @@ +# 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, AsyncGenerator +from elasticsearch7 import ( + Elasticsearch, + Transport, + RequestsHttpConnection, + ConnectionPool, + AsyncElasticsearch, + AsyncTransport, + AIOHttpConnection, +) +from elasticsearch7.helpers import ( + scan, + streaming_bulk, + reindex, + bulk, + async_scan, + async_streaming_bulk, + async_reindex, + async_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}, + ) + + +es2 = AsyncElasticsearch( + [{"host": "localhost", "port": 9443}], + transport_class=AsyncTransport, +) +t2 = 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( + es2, + query={"query": {"match_all": {}}}, + request_timeout=10, + clear_scroll=True, + scroll_kwargs={"request_timeout": 10}, + ): + pass + async for _ in async_scan( + es2, + 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(es2, async_gen()): + pass + async for _ in async_streaming_bulk(es2, async_gen().__aiter__()): + pass + async for _ in async_streaming_bulk(es2, [{}]): + pass + async for _ in async_streaming_bulk(es2, ({},)): + pass + + +async def async_bulk_types() -> None: + _, _ = await async_bulk(es2, async_gen()) + _, _ = await async_bulk(es2, async_gen().__aiter__()) + _, _ = await async_bulk(es2, [{}]) + _, _ = await async_bulk(es2, ({},)) + + +async def async_reindex_types() -> None: + _, _ = await async_reindex( + es2, "src-index", "target-index", query={"query": {"match": {"key": "val"}}} + ) + _, _ = await async_reindex( + es2, source_index="src-index", target_index="target-index", target_client=es2 + ) + _, _ = await async_reindex( + es2, + "src-index", + "target-index", + chunk_size=1, + scroll="10m", + scan_kwargs={"request_timeout": 10}, + bulk_kwargs={"request_timeout": 10}, + ) diff --git a/utils/build_dists.py b/utils/build_dists.py index b91ea187..7a31c31b 100644 --- a/utils/build_dists.py +++ b/utils/build_dists.py @@ -66,7 +66,7 @@ def test_dist(dist): # Build the venv and install the dist run("python", "-m", "venv", os.path.join(tmp_dir, "venv")) venv_python = os.path.join(tmp_dir, "venv/bin/python") - run(venv_python, "-m", "pip", "install", "-U", "pip") + run(venv_python, "-m", "pip", "install", "-U", "pip", "mypy") run(venv_python, "-m", "pip", "install", dist) # Test the sync namespaces @@ -106,6 +106,17 @@ def test_dist(dist): f"from {dist_name}.helpers import async_scan, async_bulk, async_streaming_bulk, async_reindex", ) + # Only need to test 'async_types' for non-aliased package + # since 'aliased_types' tests both async and sync. + if dist_name == "elasticsearch": + run( + venv_python, + "-m", + "mypy", + "--strict", + os.path.join(base_dir, "test_elasticsearch/test_types/async_types.py"), + ) + # Ensure that the namespaces are correct for the dist for suffix in ("", "1", "2", "5", "6", "7", "8", "9", "10"): distx_name = f"elasticsearch{suffix}" @@ -116,6 +127,27 @@ def test_dist(dist): expect_exit_code=256 if distx_name != dist_name else 0, ) + # Check that sync types work for 'elasticsearch' and + # that aliased types work for 'elasticsearchX' + if dist_name == "elasticsearch": + run( + venv_python, + "-m", + "mypy", + "--strict", + os.path.join(base_dir, "test_elasticsearch/test_types/sync_types.py"), + ) + else: + run( + venv_python, + "-m", + "mypy", + "--strict", + os.path.join( + base_dir, "test_elasticsearch/test_types/aliased_types.py" + ), + ) + # Uninstall the dist, see that we can't import things anymore run(venv_python, "-m", "pip", "uninstall", "--yes", dist_name) run( diff --git a/utils/templates/base b/utils/templates/base index 103908cc..82f5fae2 100644 --- a/utils/templates/base +++ b/utils/templates/base @@ -6,6 +6,7 @@ {{ api.description|replace("\n", " ")|wordwrap(wrapstring="\n ") }} {% endif %} {% if api.doc_url %} + `<{{ api.doc_url }}>`_ {% endif %} {% if api.params|list|length %} diff --git a/utils/templates/func_params_pyi b/utils/templates/func_params_pyi index b35a32e5..3591623a 100644 --- a/utils/templates/func_params_pyi +++ b/utils/templates/func_params_pyi @@ -2,6 +2,8 @@ {% if info.required %}{{ p }}: {{ info.type }}, {% endif %} {% endfor %} +*, + {% if api.body %} body{% if not api.body.required %}: Optional[Any]=...{% else %}: Any{% endif %}, {% endif %}