update

Author: hsz1273327

.github/workflows/python-package.yml

@@ -5,7 +5,7 @@ name: Python package
 
 on:
   push:
-    branches: [master,dev]
+    branches: [master, dev]
   pull_request:
     branches: [master]
 
@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.8, 3.9]
 
     steps:
       - uses: actions/checkout@v2
@@ -32,3 +32,6 @@ jobs:
       - name: Lint with pep8
         run: |
           pycodestyle --max-line-length=140 --ignore=E501 --first --statistics postgresql_kernel
+      - name: Static check with mypy
+        run: |
+          mypy --ignore-missing-imports --show-column-numbers --follow-imports=silent --check-untyped-defs --disallow-untyped-defs --no-implicit-optional --warn-unused-ignores --scripts-are-modules postgresql_kernel

CHANGELOG.md

@@ -1,3 +1,9 @@
+# 0.0.1
+
+1. 增加类型注释
+2. 更新依赖项,将psycopg2改为psycopg2-binary
+3. 不再支持python3.7及以下版本
+
 # 0.0.0
 
-项目创建,原样复刻了[postgres_kernel](https://github.com/bgschiller/postgres_kernel)的功能
+项目创建,原样复刻了[postgres_kernel](https://github.com/bgschiller/postgres_kernel)的功.

postgresql_kernel/kernel.py

@@ -1,4 +1,5 @@
 import os
+from pydoc import Helper
 import re
 import sys
 import json
@@ -14,16 +15,46 @@
 import psycopg2
 from psycopg2 import Error, OperationalError
 from psycopg2.extensions import (
-    QueryCanceledError, POLL_OK, POLL_READ, POLL_WRITE, STATUS_BEGIN,
+    QueryCanceledError, POLL_OK, POLL_READ, POLL_WRITE, STATUS_BEGIN, Union
 )
-
+from typing import Any, List, Optional, TypedDict, Iterable, Sequence, Mapping, Union, Tuple
 from . import __version__
 
+KernelDictType = TypedDict(
+    'KernelDictType',
+    {
+        "argv": List[str],
+        "display_name": str,
+        "mimetype": str,
+        "language": str,
+        "name": str
+    },
+    total=False)
+
+HelpLinkDictType = TypedDict(
+    'HelpLinkDictType',
+    {
+        "text": str,
+        "url": str
+    },
+    total=False)
+
+LanguageInfoDictType = TypedDict(
+    'LanguageInfoDictType',
+    {
+        'mimetype': str,
+        'name': str,
+        'file_extension': str,
+        'version': str,
+        'help_links': List[HelpLinkDictType]
+    },
+    total=False)
+
 version_pat = re.compile(r'^PostgreSQL (\d+(\.\d+)+)')
 CONN_STRING_COMMENT = re.compile(r'--\s*connection:\s*(.*)\s*$')
 AUTOCOMMIT_SWITCH_COMMENT = re.compile(r'--\s*autocommit:\s*(\w+)\s*$')
 
-HELP_LINKS = [
+HELP_LINKS: List[HelpLinkDictType] = [
     {
         'text': "PostgreSQL",
         'url': "http://www.postgres.cn/docs/12/index.html",
@@ -32,15 +63,12 @@
         'text': "SQL",
         'url': "https://blog.hszofficial.site/TutorialForSQL/#/",
     },
-    # {
-    #     'text': "PostgreSQL Kernel",
-    #     'url': "https://github.com/Calysto/octave_kernel",
-    # },
+
 
 ] + MetaKernel.help_links
 
 
-def get_kernel_json():
+def get_kernel_json() -> KernelDictType:
     """Get the kernel json for the kernel."""
     here = Path(__file__)
     default_json_file = here.parent.joinpath('kernel.json')
@@ -51,11 +79,11 @@ def get_kernel_json():
     return data
 
 
-def wait_select_inter(conn):
-    """等待连接建立
+def wait_select_inter(conn: psycopg2.connection) -> None:
+    """等待连接建立.
 
     Args:
-        conn (psycopg2.Connection): pg的连接
+        conn (psycopg2.connection): pg的连接
 
     Raises:
         conn.OperationalError: 连接报错
@@ -83,7 +111,7 @@ class MissingConnection(Exception):
 
 
 class RowsDisplay:
-    def __init__(self, header, rows):
+    def __init__(self, header: Union[str, Mapping[str, str], Sequence[str]], rows: Union[Mapping[str, Iterable], Iterable[Iterable]]) -> None:
         self.header = header
         self.rows = rows
 
@@ -98,44 +126,53 @@ def _repr_latex_(self) -> str:
 
 
 class PostgreSQLKernel(MetaKernel):
-    app_name = 'postgresql_kernel'
-    implementation = 'PostgreSQL Kernel'
-    implementation_version = __version__,
-    language = 'sql'
-    help_links = HELP_LINKS
-    kernel_json = Dict(get_kernel_json()).tag(config=True)
-    _language_version = None
-    _banner = None
+
+    app_name: str = 'postgresql_kernel'
+    implementation: str = 'PostgreSQL Kernel'
+    implementation_version: str = __version__
+    language: str = 'sql'
+    help_links: List[HelpLinkDictType] = HELP_LINKS
+    kernel_json: Dict = Dict(get_kernel_json()).tag(config=True)
+    _language_version: Optional[str] = None
+    _banner: Optional[str] = None
+
+    _conn_string: str
+    _autocommit: bool
+    _conn: Optional[psycopg2.connection]
 
     @property
-    def language_version(self):
+    def language_version(self) -> str:
         if self._language_version:
             return self._language_version
-
-        m = version_pat.search(self.banner)
-        if m:
-            self._language_version = m.group(1)
-            return self._language_version
+        if self.banner:
+            m = version_pat.search(self.banner)
+            if m:
+                self._language_version = m.group(1)
+                return self._language_version
+            else:
+                return "unknown"
         else:
             return "unknown"
 
     @property
-    def banner(self):
+    def banner(self) -> Optional[str]:
         if self._banner is None:
             if self._conn is None:
                 return 'not yet connected to a database'
-            self._banner = self.fetchone('SELECT VERSION();')[0]
+            res = self.fetchone('SELECT VERSION();')
+            if res and len(res) >= 1:
+                self._banner = res[0]
         return self._banner
 
     @property
-    def language_info(self):
+    def language_info(self) -> LanguageInfoDictType:
         return {'mimetype': 'text/x-sql',
                 'name': 'sql',
                 'file_extension': '.sql',
                 'version': self.language_version,
                 'help_links': HELP_LINKS}
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         super(PostgreSQLKernel, self).__init__(*args, **kwargs)
         psycopg2.extensions.set_wait_callback(wait_select_inter)
         self._conn_string = os.getenv('DATABASE_URL', '')
@@ -144,7 +181,7 @@ def __init__(self, *args, **kwargs):
         if self._conn_string:
             self._start_connection()
 
-    def _start_connection(self):
+    def _start_connection(self) -> None:
         """与pg建立连接."""
         self.log.info('starting connection')
         try:
@@ -156,23 +193,25 @@ def _start_connection(self):
             self.send_response(self.iopub_socket, 'stream',
                                {'name': 'stderr', 'text': message})
 
-    def fetchone(self, query):
+    def fetchone(self, query: str) -> Optional[Sequence[Union[str, bytes, int, float]]]:
         """拉取一行数据
 
         Args:
             query (str): 请求的sql语句
 
         Returns:
-            [type]: [description]
+            Optional[Sequence[Union[str, bytes, int, float]]]: 获取的结果
         """
-        self.log.info(f'fetching one from: \n{query}')
-        with self._conn.cursor() as c:
-            c.execute(query)
-            one = c.fetchone()
-            self.log.info(one)
-            return one
-
-    def fetchall(self, query):
+        if self._conn:
+            self.log.info(f'fetching one from: \n{query}')
+            with self._conn.cursor() as c:
+                c.execute(query)
+                one = c.fetchone()
+                self.log.info(one)
+                return one
+        return None
+
+    def fetchall(self, query: str) -> Tuple[Optional[List[str]], Optional[List[Sequence[Union[str, bytes, int, float]]]]]:
         """拉取多行数据.
 
         Args:
@@ -181,21 +220,24 @@ def fetchall(self, query):
         Returns:
             [type]: [description]
         """
-        self.log.info(f'fetching all from: \n{query}')
-        with self._conn.cursor() as c:
-            c.execute(query)
-            desc = c.description
-            if desc:
-                keys = [col[0] for col in desc]
-                return keys, c.fetchall()
-            return None, None
-
-    def change_connection(self, conn_string):
+        if not self._conn:
+            raise Exception("need to connect to pg first")
+        else:
+            self.log.info(f'fetching all from: \n{query}')
+            with self._conn.cursor() as c:
+                c.execute(query)
+                desc = c.description
+                if desc:
+                    keys = [col[0] for col in desc]
+                    return keys, c.fetchall()
+                return None, None
+
+    def change_connection(self, conn_string: str) -> None:
         """更换连接的库."""
         self._conn_string = conn_string
         self._start_connection()
 
-    def switch_autocommit(self, switch_to):
+    def switch_autocommit(self, switch_to: bool) -> bool:
         """切换是否要自动提交."""
         self._autocommit = switch_to
         committed = False
@@ -208,7 +250,7 @@ def switch_autocommit(self, switch_to):
             self._start_connection()
         return committed
 
-    def change_autocommit_mode(self, switch):
+    def change_autocommit_mode(self, switch: str) -> None:
         """根据输入的字符串切换是否要自动提交.
 
         如果输入的字符串的全小写是true或者false则按指定的值设置,否则抛出错误.
@@ -223,7 +265,7 @@ def change_autocommit_mode(self, switch):
             )
 
         switch_bool = (parsed_switch == 'true')
-        committed = self.switch_autocommit(switch_bool) | ''
+        committed = "True" if self.switch_autocommit(switch_bool) else ''
         message = f'committed current transaction & {committed} switched autocommit mode to {self._autocommit}'
 
         self.send_response(
@@ -233,7 +275,7 @@ def change_autocommit_mode(self, switch):
             }
         )
 
-    def get_kernel_help_on(self, info, level=1, none_on_fail=False):
+    def get_kernel_help_on(self, info: Mapping[str, str], level: int = 1, none_on_fail: bool = False) -> Optional[str]:
         self.log.warning("get kernel help")
         code = info['code'].strip()
         if not code or len(code.split()) > 1:
@@ -244,7 +286,7 @@ def get_kernel_help_on(self, info, level=1, none_on_fail=False):
         shell_magic = self.line_magics['shell']
         return shell_magic.get_help_on(info, 1)
 
-    def do_execute_meta(self, code):
+    def do_execute_meta(self, code: str) -> ExceptionWrapper:
         """
         Execute meta code in the kernel. This uses the execute infrastructure
         but allows JavaScript to talk directly to the kernel bypassing normal
@@ -261,11 +303,11 @@ def do_execute_meta(self, code):
             tb = traceback.format_exception(exc_type, exc_value, exc_traceback)
             return ExceptionWrapper(ename=str(type(e)), evalue=str(e), traceback=tb)
 
-    def do_execute_direct(self, code, silent=False):
+    def do_execute_direct(self, code: str, silent: bool = False) -> Optional[Union[RowsDisplay, ExceptionWrapper]]:
         if code.strip().lower() in ['quit', 'quit()', 'exit', 'exit()']:
             self.do_shutdown(True)
             self.payload = [{"source": "ask_exit"}]
-            return
+            return None
         try:
             connection_string = CONN_STRING_COMMENT.findall(code)
             autocommit_switch = AUTOCOMMIT_SWITCH_COMMENT.findall(code)
@@ -276,7 +318,7 @@ def do_execute_direct(self, code, silent=False):
 
             code = AUTOCOMMIT_SWITCH_COMMENT.sub('', CONN_STRING_COMMENT.sub('', code))
             if not code.strip():
-                return
+                return None
             if self._conn is None:
                 raise MissingConnection(f'''\
     Error: Unable to connect to a database at "{self._conn_string}".
@@ -289,8 +331,7 @@ def do_execute_direct(self, code, silent=False):
                 self._conn.rollback()
                 raise qce
             except Error as e:
-                self.send_response(self.iopub_socket, 'stream',
-                                   {'name': 'stderr', 'text': str(e)})
+                self.send_response(self.iopub_socket, 'stream', {'name': 'stderr', 'text': str(e)})
                 self._conn.rollback()
                 raise e
             else:
@@ -308,9 +349,10 @@ def do_execute_direct(self, code, silent=False):
                             'text': str(notice)
                         })
                 self._conn.notices = []
-                if header is not None and len(rows) > 0:
+                if header is not None and rows and len(rows) > 0:
                     return RowsDisplay(header, rows)
-
+                else:
+                    return None
         except Exception as e:
             exc_type, exc_value, exc_traceback = sys.exc_info()
             tb = traceback.format_exception(exc_type, exc_value, exc_traceback)