feat: type annotations, consolidate on pytest (#13)

Author: barakalon

.github/workflows/tests.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.7", "3.8", "3.9", "3.10"]
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
@@ -29,3 +29,6 @@ jobs:
     - name: Format
       run: |
         make format-check
+    - name: Type annotations
+      run: |
+        make types

Makefile

@@ -17,11 +17,16 @@ run:
 lint:
 	python -m pylint mysql_mimic/ tests/
 
+types:
+	python -m mypy -p mysql_mimic -p tests
+
 test:
 	coverage run --source=mysql_mimic -m pytest
 	coverage report
 	coverage html
 
+check: format-check lint types test
+
 build: clean
 	python setup.py sdist bdist_wheel
 

mysql_mimic/admin.py

@@ -1,20 +1,18 @@
 import re
+from typing import Optional, Any
 
+from mysql_mimic.session import Session
 from mysql_mimic.charset import CharacterSet
 from mysql_mimic.errors import MysqlError, ErrorCode
 from mysql_mimic.results import ResultColumn, ResultSet, Column
 from mysql_mimic.types import ColumnType
+from mysql_mimic.variables import SystemVariables
 
 
 class Admin:
     """
     Administration Statements (plus some other things)
     https://dev.mysql.com/doc/refman/8.0/en/sql-server-administration-statements.html
-
-    Args:
-        connection_id (int): connection ID
-        session (mysql_mimic.session.Session): session
-        variables (mysql_mimic.variables.SystemVariables): system variables
     """
 
     # Regex for finding "information functions" to replace in SQL statements
@@ -139,18 +137,20 @@ class Admin:
         re.IGNORECASE | re.VERBOSE,
     )
 
-    def __init__(self, connection_id, session, variables):
+    def __init__(
+        self, connection_id: int, session: Session, variables: SystemVariables
+    ):
         self.connection_id = connection_id
         self.session = session
-        self.database = None
-        self.username = None
+        self.database: Optional[str] = None
+        self.username: Optional[str] = None
         self.vars = variables
 
-    def replace_variables(self, sql):
+    def replace_variables(self, sql: str) -> str:
         sql = self.REGEX_INFO_FUNC.sub(self._replace_info_func, sql)
         return self.REGEX_SESSION_VAR.sub(self._replace_session_var, sql)
 
-    def _replace_info_func(self, matchobj):
+    def _replace_info_func(self, matchobj: re.Match) -> str:
         func = (matchobj.group("func") or matchobj.group("current_user")).lower()
         if func == "connection_id":
             return str(self.connection_id)
@@ -164,13 +164,13 @@ def _replace_info_func(self, matchobj):
             return f"'{self.database}'" if self.database else "NULL"
         raise MysqlError(f"Failed to parse system information function: {func}")
 
-    def _replace_session_var(self, matchobj):
+    def _replace_session_var(self, matchobj: re.Match) -> str:
         var = matchobj.group("var").lower()
         if var in self.vars:
             return f"'{self.vars[var]}'"
         raise MysqlError(f"Unknown variable: {var}", ErrorCode.UNKNOWN_SYSTEM_VARIABLE)
 
-    async def parse(self, sql):
+    async def parse(self, sql: str) -> Optional[ResultSet]:
         m = self.REGEX_CMD.match(sql)
         if not m:
             return None
@@ -187,7 +187,7 @@ async def parse(self, sql):
             return ResultSet([], [])
         raise MysqlError("Failed to parse command", ErrorCode.PARSE_ERROR)
 
-    async def _parse_set(self, this):
+    async def _parse_set(self, this: str) -> Optional[ResultSet]:
         m = self.REGEX_SET_NAMES.match(this)
         if m:
             return await self._parse_set_names(
@@ -209,11 +209,13 @@ async def _parse_set(self, this):
 
         raise MysqlError("Unsupported SET command", ErrorCode.NOT_SUPPORTED_YET)
 
-    async def _set_variable(self, key, val):
+    async def _set_variable(self, key: str, val: Any) -> None:
         self.vars[key] = val
         await self.session.set(**{key: val})
 
-    async def _parse_set_names(self, charset_name, collation_name):
+    async def _parse_set_names(
+        self, charset_name: str, collation_name: str
+    ) -> Optional[ResultSet]:
         await self._set_variable("character_set_client", charset_name)
         await self._set_variable("character_set_connection", charset_name)
         await self._set_variable("character_set_results", charset_name)
@@ -224,14 +226,25 @@ async def _parse_set_names(self, charset_name, collation_name):
         return ResultSet([], [])
 
     async def _parse_set_variables(
-        self, global_, persist, persist_only, session, user, var_name, value
-    ):  # pylint: disable=unused-argument
+        self,
+        global_: Optional[str],
+        persist: Optional[str],
+        persist_only: Optional[str],
+        session: Optional[str],
+        user: Optional[str],
+        var_name: Optional[str],
+        value: Optional[str],
+    ) -> Optional[ResultSet]:  # pylint: disable=unused-argument
         if global_ or persist or persist_only or user:
             raise MysqlError(
                 "Only setting session variables is supported",
                 ErrorCode.NOT_SUPPORTED_YET,
             )
 
+        assert value is not None
+        assert var_name is not None
+        result: Any
+
         value_lower = value.lower()
         if value_lower in {"true", "on"}:
             result = True
@@ -256,7 +269,7 @@ async def _parse_set_variables(
         await self._set_variable(var_name, result)
         return ResultSet([], [])
 
-    async def _parse_show(self, this):
+    async def _parse_show(self, this: str) -> Optional[ResultSet]:
         m = self.REGEX_SHOW_VARS.match(this)
         if m:
             return await self._parse_show_variables(
@@ -284,22 +297,29 @@ async def _parse_show(self, this):
 
         raise MysqlError("Unsupported SHOW command", ErrorCode.NOT_SUPPORTED_YET)
 
-    def _like_to_regex(self, like):
+    def _like_to_regex(self, like: Optional[str]) -> re.Pattern:
         if like is None:
             return re.compile(r".*")
         like = like.replace("%", ".*")
         like = like.replace("_", ".")
         return re.compile(like)
 
     async def _parse_show_columns(
-        self, extended, full, db_name, tbl_name, like
-    ):  # pylint: disable=unused-argument
+        self,
+        extended: bool,
+        full: bool,
+        db_name: Optional[str],
+        tbl_name: Optional[str],
+        like: re.Pattern,
+    ) -> Optional[ResultSet]:  # pylint: disable=unused-argument
         db_name = db_name or self.database
         if not db_name:
             raise MysqlError("No database selected", ErrorCode.NO_DB_ERROR)
 
-        columns = await self.session.show_columns(db_name, tbl_name)
-        columns = [c if isinstance(c, Column) else Column(**c) for c in columns]
+        columns = [
+            c if isinstance(c, Column) else Column(**c)
+            for c in await self.session.show_columns(db_name, tbl_name or "")
+        ]
         columns = [c for c in columns if like.match(c.name)]
 
         result_columns = [
@@ -333,8 +353,8 @@ async def _parse_show_columns(
         return ResultSet(rows=rows, columns=result_columns)
 
     async def _parse_show_index(
-        self, extended, db_name, tbl_name
-    ):  # pylint: disable=unused-argument
+        self, extended: bool, db_name: Optional[str], tbl_name: Optional[str]
+    ) -> ResultSet:  # pylint: disable=unused-argument
         result_columns = [
             ResultColumn("Table", ColumnType.STRING),
             ResultColumn("Non_unique", ColumnType.TINY),
@@ -355,8 +375,8 @@ async def _parse_show_index(
         return ResultSet(rows=[], columns=result_columns)
 
     async def _parse_show_variables(
-        self, global_, session, like
-    ):  # pylint: disable=unused-argument
+        self, global_: bool, session: bool, like: re.Pattern
+    ) -> ResultSet:  # pylint: disable=unused-argument
         rows = list(self.vars.items())
         rows = [(k, v) for k, v in rows if like.match(k)]
         return ResultSet(

mysql_mimic/auth.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import io
 from copy import copy
 from hashlib import sha1
 import logging
 import secrets
 from dataclasses import dataclass
-from typing import Optional, Dict
+from typing import Optional, Dict, AsyncGenerator, Union, Tuple
 
 from mysql_mimic.types import read_str_null
 from mysql_mimic.utils import xor
@@ -35,21 +37,25 @@ class AuthInfo:
     data: bytes
     user: User
     connect_attrs: Dict[str, str]
-    client_plugin_name: str
-    handshake_auth_data: bytes
+    client_plugin_name: Optional[str]
+    handshake_auth_data: Optional[bytes]
     handshake_plugin_name: str
 
-    def copy(self, data):
+    def copy(self, data: bytes) -> AuthInfo:
         new = copy(self)
         new.data = data
         return new
 
 
+Decision = Union[Success, Forbidden, bytes]
+AuthState = AsyncGenerator[Decision, AuthInfo]
+
+
 class AuthPlugin:
     name = ""
-    client_plugin_name = None  # None means any
+    client_plugin_name: Optional[str] = None  # None means any
 
-    async def auth(self, auth_info=None):
+    async def auth(self, auth_info: Optional[AuthInfo] = None) -> AuthState:
         """
         Begin the authentication lifecycle.
 
@@ -58,7 +64,9 @@ async def auth(self, auth_info=None):
         """
         yield Forbidden()
 
-    async def start(self, auth_info=None):
+    async def start(
+        self, auth_info: Optional[AuthInfo] = None
+    ) -> Tuple[Decision, AuthState]:
         state = self.auth(auth_info)
         data = await state.__anext__()
         return data, state
@@ -67,7 +75,7 @@ async def start(self, auth_info=None):
 class GullibleAuthPlugin(AuthPlugin):
     name = "mysql_mimic_gullible"
 
-    async def auth(self, auth_info=None):
+    async def auth(self, auth_info: Optional[AuthInfo] = None) -> AuthState:
         if not auth_info:
             auth_info = yield b"\x00" * 20  # 20 bytes of filler to be ignored
         yield Success(authenticated_as=auth_info.username)
@@ -77,27 +85,27 @@ class AbstractMysqlClearPasswordAuthPlugin(AuthPlugin):
     name = "abstract_mysql_clear_password"
     client_plugin_name = "mysql_clear_password"
 
-    async def auth(self, auth_info=None):
+    async def auth(self, auth_info: Optional[AuthInfo] = None) -> AuthState:
         if not auth_info:
             auth_info = yield b"\x00" * 20  # 20 bytes of filler to be ignored
 
         r = io.BytesIO(auth_info.data)
         password = read_str_null(r).decode()
         authenticated_as = await self.check(auth_info.username, password)
-        if authenticated_as:
+        if authenticated_as is not None:
             yield Success(authenticated_as)
         else:
             yield Forbidden()
 
-    async def check(self, username, password):
+    async def check(self, username: str, password: str) -> Optional[str]:
         return username
 
 
 class MysqlNativePasswordAuthPlugin(AuthPlugin):
     name = "mysql_native_password"
     client_plugin_name = "mysql_native_password"
 
-    async def auth(self, auth_info=None):
+    async def auth(self, auth_info: Optional[AuthInfo] = None) -> AuthState:
         if (
             auth_info
             and auth_info.handshake_plugin_name == self.name
@@ -138,5 +146,5 @@ async def auth(self, auth_info=None):
             yield Forbidden()
 
 
-def get_mysql_native_password_auth_string(password):
+def get_mysql_native_password_auth_string(password: str) -> str:
     return sha1(sha1(password.encode("utf-8")).digest()).hexdigest()

mysql_mimic/charset.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from enum import IntEnum
 
 
@@ -45,19 +46,19 @@ class CharacterSet(IntEnum):
     utf8mb4 = 255
 
     @property
-    def codec(self):
+    def codec(self) -> str:
         if self.name == "utf8mb4":
             return "utf8"
         return self.name
 
     @property
-    def default_collation(self):
+    def default_collation(self) -> Collation:
         return DEFAULT_COLLATIONS[self]
 
-    def decode(self, b):
+    def decode(self, b: bytes) -> str:
         return b.decode(self.codec)
 
-    def encode(self, s):
+    def encode(self, s: str) -> bytes:
         return s.encode(self.codec)
 
 
@@ -287,11 +288,11 @@ class Collation(IntEnum):
     utf8mb4_0900_ai_ci = 255
 
     @property
-    def codec(self):
+    def codec(self) -> str:
         return self.charset.codec
 
     @property
-    def charset(self):
+    def charset(self) -> CharacterSet:
         return DEFAULT_CHARACTER_SETS[self]