Basic mypy support

.github/workflows/style.yaml

@@ -12,8 +12,8 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: [3.8]
-        environment: [style, black]
+        python-version: [3.9]
+        environment: [style, black, mypy]
 
     steps:
     - uses: actions/checkout@v1

pyproject.toml

@@ -1,2 +1,127 @@
 [tool.black]
 line-length = 120
+
+[tool.mypy]
+python_version = "3.9"
+warn_return_any = false
+warn_unused_configs = true
+warn_redundant_casts = true
+warn_unused_ignores = false
+strict_equality = false
+check_untyped_defs = false
+disallow_untyped_decorators = false
+no_implicit_optional = false
+show_error_codes = true
+
+# Start with more lenient settings, can be made stricter over time
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+
+exclude = [
+    "ydb/_grpc/v3/",
+    "ydb/_grpc/v4/",
+    "ydb/_grpc/v5/",
+    "ydb/_grpc/v6/",
+    "tests/",
+    "examples/",
+    "docs/",
+    # Exclude test files inside ydb directory
+    ".*_test\\.py$",
+]
+
+[[tool.mypy.overrides]]
+module = [
+    "ydb._grpc.v3.*",
+    "ydb._grpc.v4.*",
+    "ydb._grpc.v5.*",
+    "ydb._grpc.v6.*",
+    "ydb._grpc.grpcwrapper.*",
+    "ydb._grpc.common.*",
+    "google.protobuf.*",
+    "grpc.*",
+    "aiohttp.*",
+    "yandexcloud.*",
+    "yandex.cloud.*",
+    "jwt.*",
+    "kikimr.*",
+    "sqlalchemy.*",
+    "requests.*",
+    "ydb.public.api.*",
+    "contrib.ydb.public.api.*",
+]
+ignore_missing_imports = true
+ignore_errors = true
+
+# Files that heavily use protobuf types - allow attr-defined errors
+[[tool.mypy.overrides]]
+module = [
+    "ydb.types",
+    "ydb.convert",
+    "ydb._apis",
+    "ydb._utilities",
+]
+disable_error_code = ["attr-defined", "union-attr", "arg-type", "return-value", "assignment", "misc"]
+
+# IAM and OAuth modules with optional imports
+[[tool.mypy.overrides]]
+module = [
+    "ydb.iam.*",
+    "ydb.oauth2_token_exchange.*",
+]
+disable_error_code = ["assignment", "arg-type", "return-value"]
+
+# Query modules - now fully typed!
+
+# Table modules - protobuf attr-defined issues
+[[tool.mypy.overrides]]
+module = ["ydb.table", "ydb.aio.table"]
+disable_error_code = ["attr-defined", "return-value", "arg-type"]
+
+# Topic modules
+[[tool.mypy.overrides]]
+module = [
+    "ydb._topic_reader.*",
+    "ydb._topic_writer.*",
+    "ydb._topic_common.*",
+]
+disable_error_code = [
+    "arg-type",
+    "return-value",
+    "assignment",
+    "misc",
+    "empty-body",
+    "var-annotated",
+    "dict-item",
+    "list-item",
+    "no-redef",
+    "union-attr",
+    "index",
+    "truthy-function",
+    "call-arg",
+    "attr-defined",
+    "return",
+]
+
+
+# Async IAM and OAuth
+[[tool.mypy.overrides]]
+module = [
+    "ydb.aio.iam",
+    "ydb.aio.oauth2_token_exchange",
+    "ydb.aio.credentials",
+]
+disable_error_code = ["assignment", "override"]
+
+# Topic module
+[[tool.mypy.overrides]]
+module = [
+    "ydb.topic",
+]
+disable_error_code = ["assignment", "attr-defined"]
+
+# SQLAlchemy integration
+[[tool.mypy.overrides]]
+module = [
+    "ydb.sqlalchemy.*",
+]
+disable_error_code = ["no-redef"]

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py,py-proto4,py-proto3,py-tls,py-tls-proto4,py-tls-proto3,style,pylint,black,protoc,py-cov-proto4
+envlist = py,py-proto6,py-proto5,py-proto4,py-proto3,style,pylint,black,black-format,mypy
 minversion = 4.2.6
 skipsdist = True
 ignore_basepython_conflict = true
@@ -63,6 +63,16 @@ deps = pylint
     -r{toxinidir}/test-requirements.txt
 commands = pylint ydb
 
+[testenv:mypy]
+skip_install = false
+deps =
+    mypy>=1.0.0
+    types-protobuf
+    types-aiofiles
+    types-requests
+    -r{toxinidir}/requirements.txt
+commands = mypy ydb
+
 [testenv:style]
 commands =
    pytest --flake8 -m flake8

ydb/_grpc/grpcwrapper/common_utils.py

@@ -25,8 +25,7 @@
 from google.protobuf.duration_pb2 import Duration as ProtoDuration
 from google.protobuf.timestamp_pb2 import Timestamp as ProtoTimeStamp
 
-from ...driver import Driver
-from ...aio.driver import Driver as DriverIO
+from ..._typing import SupportedDriverType
 
 # Workaround for good IDE and universal for runtime
 if typing.TYPE_CHECKING:
@@ -148,7 +147,7 @@ def close(self):
         ...
 
 
-SupportedDriverType = Union[Driver, DriverIO]
+# SupportedDriverType imported from ydb._typing
 
 
 class GrpcWrapperAsyncIO(IGrpcWrapperAsyncIO):
@@ -194,7 +193,7 @@ def _clean_executor(self, wait: bool):
         if self._wait_executor:
             self._wait_executor.shutdown(wait)
 
-    async def _start_asyncio_driver(self, driver: DriverIO, stub, method):
+    async def _start_asyncio_driver(self, driver, stub, method):
         requests_iterator = QueueToIteratorAsyncIO(self.from_client_grpc)
         stream_call = await driver(
             requests_iterator,
@@ -205,7 +204,7 @@ async def _start_asyncio_driver(self, driver: DriverIO, stub, method):
         self._stream_call = stream_call
         self.from_server_grpc = stream_call.__aiter__()
 
-    async def _start_sync_driver(self, driver: Driver, stub, method):
+    async def _start_sync_driver(self, driver, stub, method):
         requests_iterator = AsyncQueueToSyncIteratorAsyncIO(self.from_client_grpc)
         self._wait_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -164,7 +164,7 @@ def to_proto(self) -> ydb_query_pb2.TransactionControl:
 class ExecuteQueryRequest(IToProto):
     session_id: str
     query_content: QueryContent
-    tx_control: TransactionControl
+    tx_control: Optional[TransactionControl]
     concurrent_result_sets: bool
     exec_mode: int
     parameters: dict

ydb/_typing.py

@@ -0,0 +1,86 @@
+"""
+Common type definitions for YDB Python SDK.
+
+This module contains type aliases, protocols, and type variables
+used across the SDK for proper typing support.
+"""
+
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Protocol,
+    Tuple,
+    TypeVar,
+    Union,
+    TYPE_CHECKING,
+)
+
+import grpc
+
+if TYPE_CHECKING:
+    from .driver import Driver as _SyncDriver
+    from .aio.driver import Driver as _AsyncDriver
+
+
+# =============================================================================
+# Driver Type Variables
+# =============================================================================
+
+# TypeVar constrained to sync or async driver - use in Generic classes
+# Example: class BaseQuerySession(Generic[DriverT]): ...
+DriverT = TypeVar("DriverT", "_SyncDriver", "_AsyncDriver")
+
+# Union type for functions that accept either driver
+SupportedDriverType = Union["_SyncDriver", "_AsyncDriver"]
+
+
+# =============================================================================
+# gRPC Stream Types
+# =============================================================================
+
+# gRPC streaming calls return an object that is both grpc.Call (with cancel())
+# and an Iterator. Since grpc doesn't export a public type for this combination,
+# we define a type that matches the actual runtime behavior.
+# See: grpc._channel._MultiThreadedRendezvous which inherits from grpc.Call, grpc.Future
+_StreamItemT = TypeVar("_StreamItemT", covariant=True)
+
+
+class GrpcStreamCall(grpc.Call, Iterable[_StreamItemT]):
+    """Type for gRPC streaming call response.
+
+    gRPC streaming calls return _MultiThreadedRendezvous which is both
+    a grpc.Call (with cancel()) and an Iterator. This class provides
+    proper typing by inheriting from both.
+
+    Usage:
+        _stream: Optional[GrpcStreamCall[SessionState]] = None
+    """
+
+    pass
+
+
+# =============================================================================
+# RPC Call Signatures
+# =============================================================================
+
+# Type for wrap_result callback
+WrapResultFunc = Callable[..., Any]
+
+# Type for RPC call arguments tuple
+WrapArgsType = Tuple[Any, ...]
+
+
+# =============================================================================
+# Lock Protocol
+# =============================================================================
+
+
+class LockProtocol(Protocol):
+    """Protocol for lock objects - supports both threading.Lock and fake locks."""
+
+    def __enter__(self) -> Any:
+        ...
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        ...