From 0ea73f1fb0f6aee69ec8ec0b1971f1580071f3d6 Mon Sep 17 00:00:00 2001
From: Timofey Koolin <timofey.koolin@gmail.com>
Date: Mon, 6 Feb 2023 11:46:29 +0300
Subject: [PATCH] improve grpc import for friendly for IDE

---
 ydb/_apis.py         | 53 +++++++++++++++++++++++++++++++-------------
 ydb/credentials.py   | 11 +++++++--
 ydb/export.py        | 12 ++++++++--
 ydb/import_client.py | 13 +++++++++--
 ydb/scripting.py     | 12 ++++++++--
 5 files changed, 77 insertions(+), 24 deletions(-)

diff --git a/ydb/_apis.py b/ydb/_apis.py
index 8ee2d731..6f2fc3ab 100644
--- a/ydb/_apis.py
+++ b/ydb/_apis.py
@@ -1,21 +1,42 @@
 # -*- coding: utf-8 -*-
-from ydb._grpc.common import (
-    ydb_cms_v1_pb2_grpc,
-    ydb_discovery_v1_pb2_grpc,
-    ydb_scheme_v1_pb2_grpc,
-    ydb_table_v1_pb2_grpc,
-)
-from ydb._grpc.common.protos import (
-    ydb_status_codes_pb2,
-    ydb_discovery_pb2,
-    ydb_scheme_pb2,
-    ydb_table_pb2,
-    ydb_value_pb2,
-)
-from ydb._grpc.common.protos import ydb_operation_pb2
-from ydb._grpc.common.protos import ydb_common_pb2
-from ydb._grpc.common import ydb_operation_v1_pb2_grpc
+# Workaround for good IDE and universal for runtime
+# noinspection PyUnreachableCode
+if False:
+    from ._grpc.v4 import (
+        ydb_cms_v1_pb2_grpc,
+        ydb_discovery_v1_pb2_grpc,
+        ydb_scheme_v1_pb2_grpc,
+        ydb_table_v1_pb2_grpc,
+        ydb_operation_v1_pb2_grpc,
+    )
 
+    from ._grpc.v4.protos import (
+        ydb_status_codes_pb2,
+        ydb_discovery_pb2,
+        ydb_scheme_pb2,
+        ydb_table_pb2,
+        ydb_value_pb2,
+        ydb_operation_pb2,
+        ydb_common_pb2,
+    )
+else:
+    from ._grpc.common import (
+        ydb_cms_v1_pb2_grpc,
+        ydb_discovery_v1_pb2_grpc,
+        ydb_scheme_v1_pb2_grpc,
+        ydb_table_v1_pb2_grpc,
+        ydb_operation_v1_pb2_grpc,
+    )
+
+    from ._grpc.common.protos import (
+        ydb_status_codes_pb2,
+        ydb_discovery_pb2,
+        ydb_scheme_pb2,
+        ydb_table_pb2,
+        ydb_value_pb2,
+        ydb_operation_pb2,
+        ydb_common_pb2,
+    )
 
 StatusIds = ydb_status_codes_pb2.StatusIds
 FeatureFlag = ydb_common_pb2.FeatureFlag
diff --git a/ydb/credentials.py b/ydb/credentials.py
index d3c0d3ca..8e22fe2a 100644
--- a/ydb/credentials.py
+++ b/ydb/credentials.py
@@ -7,8 +7,15 @@
 from concurrent import futures
 import logging
 import time
-from ydb._grpc.common.protos import ydb_auth_pb2
-from ydb._grpc.common import ydb_auth_v1_pb2_grpc
+
+# Workaround for good IDE and universal for runtime
+# noinspection PyUnreachableCode
+if False:
+    from ._grpc.v4.protos import ydb_auth_pb2
+    from ._grpc.v4 import ydb_auth_v1_pb2_grpc
+else:
+    from ._grpc.common.protos import ydb_auth_pb2
+    from ._grpc.common import ydb_auth_v1_pb2_grpc
 
 
 YDB_AUTH_TICKET_HEADER = "x-ydb-auth-ticket"
diff --git a/ydb/export.py b/ydb/export.py
index 23df3c63..419a753b 100644
--- a/ydb/export.py
+++ b/ydb/export.py
@@ -3,8 +3,16 @@
 from . import _apis
 
 from . import settings_impl as s_impl
-from ydb._grpc.common.protos import ydb_export_pb2
-from ydb._grpc.common import ydb_export_v1_pb2_grpc
+
+# Workaround for good IDE and universal for runtime
+# noinspection PyUnreachableCode
+if False:
+    from ._grpc.v4.protos import ydb_export_pb2
+    from ._grpc.v4 import ydb_export_v1_pb2_grpc
+else:
+    from ._grpc.common.protos import ydb_export_pb2
+    from ._grpc.common import ydb_export_v1_pb2_grpc
+
 from . import operation
 
 _ExportToYt = "ExportToYt"
diff --git a/ydb/import_client.py b/ydb/import_client.py
index 61199b8c..d1ccc99a 100644
--- a/ydb/import_client.py
+++ b/ydb/import_client.py
@@ -3,8 +3,17 @@
 from . import _apis
 
 from . import settings_impl as s_impl
-from ydb._grpc.common.protos import ydb_import_pb2
-from ydb._grpc.common import ydb_import_v1_pb2_grpc
+
+# Workaround for good IDE and universal for runtime
+# noinspection PyUnreachableCode
+if False:
+    from ._grpc.v4.protos import ydb_import_pb2
+    from ._grpc.v4 import ydb_import_v1_pb2_grpc
+else:
+    from ._grpc.common.protos import ydb_import_pb2
+    from ._grpc.common import ydb_import_v1_pb2_grpc
+
+
 from . import operation
 
 _ImportFromS3 = "ImportFromS3"
diff --git a/ydb/scripting.py b/ydb/scripting.py
index 247f8105..9fed037a 100644
--- a/ydb/scripting.py
+++ b/ydb/scripting.py
@@ -1,5 +1,13 @@
-from ydb._grpc.common.protos import ydb_scripting_pb2
-from ydb._grpc.common import ydb_scripting_v1_pb2_grpc
+# Workaround for good IDE and universal for runtime
+# noinspection PyUnreachableCode
+if False:
+    from ._grpc.v4.protos import ydb_scripting_pb2
+    from ._grpc.v4 import ydb_scripting_v1_pb2_grpc
+else:
+    from ._grpc.common.protos import ydb_scripting_pb2
+    from ._grpc.common import ydb_scripting_v1_pb2_grpc
+
+
 from . import issues, convert, settings