add type hints for issues.py and types.py

Author: RaileyHartheim

ydb/issues.py

@@ -2,9 +2,13 @@
 from google.protobuf import text_format
 import enum
 import queue
+import typing
 
 from . import _apis
 
+if typing.TYPE_CHECKING:
+    from _grpc.v4.protos import ydb_issue_message_pb2, ydb_operation_pb2
+
 
 _TRANSPORT_STATUSES_FIRST = 401000
 _CLIENT_STATUSES_FIRST = 402000
@@ -43,10 +47,26 @@ class StatusCode(enum.IntEnum):
     SESSION_POOL_EMPTY = _CLIENT_STATUSES_FIRST + 40
 
 
+# TODO: convert from proto IssueMessage
+class _IssueMessage:
+
+    def __init__(
+        self,
+        message: str,
+        issue_code: int,
+        severity: int,
+        issues
+    ) -> None:
+        self.message = message
+        self.issue_code = issue_code
+        self.severity = severity
+        self.issues = issues
+
+
 class Error(Exception):
     status = None
 
-    def __init__(self, message, issues=None):
+    def __init__(self, message: str, issues=None):
         super(Error, self).__init__(message)
         self.issues = issues
         self.message = message
@@ -161,14 +181,14 @@ def __init__(self, message: str):
         super().__init__(message)
 
 
-def _format_issues(issues):
+def _format_issues(issues: typing.Iterable[ydb_issue_message_pb2.IssueMessage]) -> str:
     if not issues:
         return ""
 
     return " ,".join(text_format.MessageToString(issue, as_utf8=False, as_one_line=True) for issue in issues)
 
 
-def _format_response(response):
+def _format_response(response: ydb_operation_pb2.Operation) -> str:
     fmt_issues = _format_issues(response.issues)
     return f"{fmt_issues} (server_code: {response.status})"
 
@@ -196,7 +216,7 @@ def _format_response(response):
 }
 
 
-def _process_response(response_proto):
+def _process_response(response_proto: ydb_operation_pb2.Operation) -> None:
     if response_proto.status not in _success_status_codes:
         exc_obj = _server_side_error_map.get(response_proto.status)
         raise exc_obj(_format_response(response_proto), response_proto.issues)

ydb/types.py

@@ -4,73 +4,101 @@
 import json
 from . import _utilities, _apis
 from datetime import date, datetime, timedelta
+import typing
 import uuid
 import struct
 from google.protobuf import struct_pb2
 
+from .table import TableClientSettings
+
+
+if typing.TYPE_CHECKING:
+    from ._grpc.v4.protos import ydb_value_pb2
+
 
 _SECONDS_IN_DAY = 60 * 60 * 24
 _EPOCH = datetime(1970, 1, 1)
 
 
-def _from_date(x, table_client_settings):
+def _from_date(
+        x: ydb_value_pb2.Value,
+        table_client_settings: TableClientSettings
+) -> typing.Union[date, int]:
     if table_client_settings is not None and table_client_settings._native_date_in_result_sets:
         return _EPOCH.date() + timedelta(days=x.uint32_value)
     return x.uint32_value
 
 
-def _to_date(pb, value):
+def _to_date(pb: ydb_value_pb2.Value, value: typing.Union[date, int]) -> None:
     if isinstance(value, date):
         pb.uint32_value = (value - _EPOCH.date()).days
     else:
         pb.uint32_value = value
 
 
-def _from_datetime_number(x, table_client_settings):
+def _from_datetime_number(
+        x: typing.Union[float, datetime],
+        table_client_settings: TableClientSettings
+) -> datetime:
     if table_client_settings is not None and table_client_settings._native_datetime_in_result_sets:
         return datetime.utcfromtimestamp(x)
     return x
 
 
-def _from_json(x, table_client_settings):
+def _from_json(
+        x: typing.Union[str, bytearray, bytes],
+        table_client_settings: TableClientSettings
+):
     if table_client_settings is not None and table_client_settings._native_json_in_result_sets:
         return json.loads(x)
     return x
 
 
-def _to_uuid(value_pb, table_client_settings):
+def _to_uuid(
+        value_pb: ydb_value_pb2.Value,
+        table_client_settings: TableClientSettings
+) -> uuid.UUID:
     return uuid.UUID(bytes_le=struct.pack("QQ", value_pb.low_128, value_pb.high_128))
 
 
-def _from_uuid(pb, value):
+def _from_uuid(pb: ydb_value_pb2.Value, value: uuid.UUID):
     pb.low_128 = struct.unpack("Q", value.bytes_le[0:8])[0]
     pb.high_128 = struct.unpack("Q", value.bytes_le[8:16])[0]
 
 
-def _from_interval(value_pb, table_client_settings):
+def _from_interval(
+        value_pb: ydb_value_pb2.Value,
+        table_client_settings: TableClientSettings
+) -> typing.Union[timedelta, int]:
     if table_client_settings is not None and table_client_settings._native_interval_in_result_sets:
         return timedelta(microseconds=value_pb.int64_value)
     return value_pb.int64_value
 
 
-def _timedelta_to_microseconds(value):
+def _timedelta_to_microseconds(value: timedelta) -> int:
     return (value.days * _SECONDS_IN_DAY + value.seconds) * 1000000 + value.microseconds
 
 
-def _to_interval(pb, value):
+def _to_interval(pb: ydb_value_pb2.Value, value: typing.Union[timedelta, int]) -> int:
     if isinstance(value, timedelta):
         pb.int64_value = _timedelta_to_microseconds(value)
     else:
         pb.int64_value = value
 
 
-def _from_timestamp(value_pb, table_client_settings):
+def _from_timestamp(
+        value_pb: ydb_value_pb2.Value,
+        table_client_settings: TableClientSettings
+) -> typing.Union[datetime, int]:
     if table_client_settings is not None and table_client_settings._native_timestamp_in_result_sets:
         return _EPOCH + timedelta(microseconds=value_pb.uint64_value)
     return value_pb.uint64_value
 
 
-def _to_timestamp(pb, value):
+def _to_timestamp(
+        pb: ydb_value_pb2.Value,
+        value: typing.Union[datetime, int]
+) -> int:
     if isinstance(value, datetime):
         pb.uint64_value = _timedelta_to_microseconds(value - _EPOCH)
     else:
@@ -129,13 +157,23 @@ class PrimitiveType(enum.Enum):
 
     DyNumber = _apis.primitive_types.DYNUMBER, "text_value"
 
-    def __init__(self, idn, proto_field, to_obj=None, from_obj=None):
+    def __init__(
+            self,
+            idn: ydb_value_pb2.Type.PrimitiveTypeId,
+            proto_field: typing.Optional[str],
+            to_obj=None,
+            from_obj=None
+    ):
         self._idn_ = idn
         self._to_obj = to_obj
         self._from_obj = from_obj
         self._proto_field = proto_field
 
-    def get_value(self, value_pb, table_client_settings):
+    def get_value(
+            self,
+            value_pb: ydb_value_pb2.Value,
+            table_client_settings: TableClientSettings
+    ):
         """
         Extracts value from protocol buffer
         :param value_pb: A protocol buffer
@@ -149,7 +187,7 @@ def get_value(self, value_pb, table_client_settings):
 
         return getattr(value_pb, self._proto_field)
 
-    def set_value(self, pb, value):
+    def set_value(self, pb: ydb_value_pb2.Value, value):
         """
         Sets value in a protocol buffer
         :param pb: A protocol buffer
@@ -176,7 +214,12 @@ def proto(self):
 class DataQuery(object):
     __slots__ = ("yql_text", "parameters_types", "name")
 
-    def __init__(self, query_id, parameters_types, name=None):
+    def __init__(
+            self,
+            query_id: str,
+            parameters_types: "dict[str, ydb_value_pb2.Type]",
+            name: typing.Optional[str] = None
+    ):
         self.yql_text = query_id
         self.parameters_types = parameters_types
         self.name = _utilities.get_query_hash(self.yql_text) if name is None else name
@@ -259,7 +302,10 @@ def __str__(self):
 class OptionalType(AbstractTypeBuilder):
     __slots__ = ("_repr", "_proto", "_item")
 
-    def __init__(self, optional_type):
+    def __init__(
+            self,
+            optional_type: typing.Union[AbstractTypeBuilder, PrimitiveType]
+    ):
         """
         Represents optional type that wraps inner type
         :param optional_type: An instance of an inner type
@@ -291,7 +337,10 @@ def __str__(self):
 class ListType(AbstractTypeBuilder):
     __slots__ = ("_repr", "_proto")
 
-    def __init__(self, list_type):
+    def __init__(
+            self,
+            list_type: typing.Union[AbstractTypeBuilder, PrimitiveType]
+    ):
         """
         :param list_type: List item type builder
         """
@@ -313,7 +362,11 @@ def __str__(self):
 class DictType(AbstractTypeBuilder):
     __slots__ = ("__repr", "__proto")
 
-    def __init__(self, key_type, payload_type):
+    def __init__(
+            self,
+            key_type: typing.Union[AbstractTypeBuilder, PrimitiveType],
+            payload_type: typing.Union[AbstractTypeBuilder, PrimitiveType]
+    ):
         """
         :param key_type: Key type builder
         :param payload_type: Payload type builder
@@ -341,7 +394,10 @@ def __init__(self):
         self.__elements_repr = []
         self.__proto = _apis.ydb_value.Type(tuple_type=_apis.ydb_value.TupleType())
 
-    def add_element(self, element_type):
+    def add_element(
+            self,
+            element_type: typing.Union[AbstractTypeBuilder, PrimitiveType]
+    ):
         """
         :param element_type: Adds additional element of tuple
         :return: self
@@ -366,7 +422,11 @@ def __init__(self):
         self.__members_repr = []
         self.__proto = _apis.ydb_value.Type(struct_type=_apis.ydb_value.StructType())
 
-    def add_member(self, name, member_type):
+    def add_member(
+            self,
+            name: str,
+            member_type: typing.Union[AbstractTypeBuilder, PrimitiveType]
+    ):
         """
         :param name:
         :param member_type:
@@ -393,7 +453,11 @@ def __init__(self):
         self.__columns_repr = []
         self.__proto = _apis.ydb_value.Type(struct_type=_apis.ydb_value.StructType())
 
-    def add_column(self, name, column_type):
+    def add_column(
+            self,
+            name: str,
+            column_type: typing.Union[AbstractTypeBuilder, PrimitiveType]
+    ):
         """
         :param name: A column name
         :param column_type: A column type