[7.10] Test typing metadata in build-dist script

Author: sethmlarson

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

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.

setup.py

@@ -72,7 +72,7 @@
         "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=[

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},
+    )

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(

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 %}

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 %}