Adding types

conftest.py

@@ -0,0 +1,3 @@
+def pytest_configure(config):
+    plugin = config.pluginmanager.getplugin('mypy')
+    plugin.mypy_argv.append('--check-untyped-defs')

pytest.ini

@@ -1,4 +1,5 @@
 [pytest]
 log_level=DEBUG
 log_format=%(created)f %(filename)s %(lineno)d %(levelname)s %(message)s
-log_date_format=%H:%M:%S
+log_date_format=%H:%M:%S
+#addopts=--mypy

src/pytds/login.py

@@ -68,13 +68,14 @@ def create_packet(self) -> bytes:
         from . import sspi
         import ctypes
         buf = ctypes.create_string_buffer(4096)
-        self._ctx, status, bufs = self._cred.create_context(
+        ctx, status, bufs = self._cred.create_context(
             flags=self._flags,
             byte_ordering='network',
             target_name=self._sname,
             output_buffers=[(sspi.SECBUFFER_TOKEN, buf)])
+        self._ctx = ctx
         if status == sspi.Status.SEC_I_COMPLETE_AND_CONTINUE:
-            self._ctx.complete_auth_token(bufs)
+            ctx.complete_auth_token(bufs)
         return bufs[0][1]
 
     def handle_next(self, packet: bytes) -> bytes | None:

src/pytds/smp.py

@@ -8,6 +8,9 @@
 import socket
 import errno
 from typing import Dict, Tuple
+
+from typing_extensions import Buffer
+
 from . import tds_base
 
 try:
@@ -59,7 +62,7 @@ def get_state(self) -> int | None:
     def close(self) -> None:
         self._mgr.close_smp_session(self)
 
-    def sendall(self, data: bytes, flags: int = 0) -> None:
+    def sendall(self, data: Buffer, flags: int = 0) -> None:
         self._mgr.send_packet(self, data)
 
     def _recv_internal(self, size: int) -> Tuple[int, int]:

src/pytds/tds.py

@@ -944,7 +944,7 @@ def set_state(self, state: int) -> int:
         return self.state
 
     @contextlib.contextmanager
-    def querying_context(self, packet_type) -> None:
+    def querying_context(self, packet_type: int) -> None:
         """ Context manager for querying.
 
         Sets state to TDS_QUERYING, and reverts it to TDS_IDLE if exception happens inside managed block,

src/pytds/tds_base.py

@@ -735,7 +735,7 @@ def choose_serializer(self, type_factory, collation):
         return type_factory.serializer_by_type(sql_type=self.type, collation=collation)
 
 
-class TransportProtocol(SocketType):
+class TransportProtocol(Protocol):
     def is_connected(self) -> bool:
         ...
 
@@ -749,7 +749,7 @@ class AuthProtocol(Protocol):
     def create_packet(self) -> bytes:
         ...
 
-    def handle_hext(self, packet: bytes) -> bytes | None:
+    def handle_next(self, packet: bytes) -> bytes | None:
         ...
 
     def close(self) -> None:

src/pytds/tls.py

@@ -4,6 +4,8 @@
 import socket
 from typing import Any
 
+from pytds.tds import _TdsSession
+
 try:
     import OpenSSL.SSL
     import cryptography.hazmat.backends.openssl.backend

test_requirements.txt

@@ -15,3 +15,5 @@ namedlist
 cryptography < 3.3
 sqlalchemy-pytds==1.0.0
 SQLAlchemy==2.0.22
+mypy
+pytest-mypy