Set module metadata on imported names

Author: sethmlarson

elasticsearch/__init__.py

@@ -21,6 +21,7 @@
 import re
 import warnings
 
+from ._utils import fixup_module_metadata
 from ._version import __versionstr__
 
 _version_groups = re.search(r"^(\d+)\.(\d+)\.(\d+)", __versionstr__).groups()  # type: ignore
@@ -74,3 +75,6 @@
     "UnsupportedProductError",
     "ElasticsearchWarning",
 ]
+
+fixup_module_metadata(__name__, globals())
+del fixup_module_metadata

elasticsearch/_utils.py

@@ -0,0 +1,34 @@
+#  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 re
+from typing import Any, Dict
+
+
+def fixup_module_metadata(module_name: str, namespace: Dict[str, Any]) -> None:
+    # Yoinked from python-trio/outcome, thanks Nathaniel! License: MIT
+    def fix_one(obj: Any) -> None:
+        mod = getattr(obj, "__module__", None)
+        if mod is not None and re.match(r"^elasticsearch[0-9]*\.", mod) is not None:
+            obj.__module__ = module_name
+            if isinstance(obj, type):
+                for attr_value in obj.__dict__.values():
+                    fix_one(attr_value)
+
+    for objname in namespace["__all__"]:
+        obj = namespace[objname]
+        fix_one(obj)

elasticsearch/client.py

@@ -61,6 +61,7 @@
 from ._sync.client.transform import TransformClient as TransformClient  # noqa: F401
 from ._sync.client.watcher import WatcherClient as WatcherClient  # noqa: F401
 from ._sync.client.xpack import XPackClient as XPackClient  # noqa: F401
+from ._utils import fixup_module_metadata
 
 # This file exists for backwards compatibility.
 warnings.warn(
@@ -69,3 +70,43 @@
     category=DeprecationWarning,
     stacklevel=2,
 )
+
+__all__ = [
+    "AsyncSearchClient",
+    "AutoscalingClient",
+    "CatClient",
+    "CcrClient",
+    "ClusterClient",
+    "DanglingIndicesClient",
+    "Elasticsearch",
+    "EnrichClient",
+    "EqlClient",
+    "FeaturesClient",
+    "FleetClient",
+    "GraphClient",
+    "IlmClient",
+    "IndicesClient",
+    "IngestClient",
+    "LicenseClient",
+    "LogstashClient",
+    "MigrationClient",
+    "MlClient",
+    "MonitoringClient",
+    "NodesClient",
+    "RollupClient",
+    "SearchableSnapshotsClient",
+    "SecurityClient",
+    "ShutdownClient",
+    "SlmClient",
+    "SnapshotClient",
+    "SqlClient",
+    "SslClient",
+    "TasksClient",
+    "TextStructureClient",
+    "TransformClient",
+    "WatcherClient",
+    "XPackClient",
+]
+
+fixup_module_metadata(__name__, globals())
+del fixup_module_metadata

elasticsearch/helpers/__init__.py

@@ -15,16 +15,11 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from .actions import (
-    _chunk_actions,
-    _process_bulk_chunk,
-    bulk,
-    expand_action,
-    parallel_bulk,
-    reindex,
-    scan,
-    streaming_bulk,
-)
+from .._async.helpers import async_bulk, async_reindex, async_scan, async_streaming_bulk
+from .._utils import fixup_module_metadata
+from .actions import _chunk_actions  # noqa: F401
+from .actions import _process_bulk_chunk  # noqa: F401
+from .actions import bulk, expand_action, parallel_bulk, reindex, scan, streaming_bulk
 from .errors import BulkIndexError, ScanError
 
 __all__ = [
@@ -36,19 +31,11 @@
     "parallel_bulk",
     "scan",
     "reindex",
-    "_chunk_actions",
-    "_process_bulk_chunk",
+    "async_scan",
+    "async_bulk",
+    "async_reindex",
+    "async_streaming_bulk",
 ]
 
-
-try:
-    from .._async.helpers import (
-        async_bulk,
-        async_reindex,
-        async_scan,
-        async_streaming_bulk,
-    )
-
-    __all__ += ["async_scan", "async_bulk", "async_reindex", "async_streaming_bulk"]
-except ImportError:  # pragma: nocover
-    pass
+fixup_module_metadata(__name__, globals())
+del fixup_module_metadata