Fix types and style

davidbrochart · davidbrochart · commit d2ec39eafae2 · 2025-06-10T23:08:10.000+02:00
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -20,7 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: [ '3.9', '3.10.0', '3.11', '3.12', '3.13' ]
+        python-version: [ '3.9', '3.10', '3.11', '3.12', '3.13' ]
 
     steps:
     - name: Checkout
@@ -37,8 +37,8 @@ jobs:
 
     - name: Check style and types
       run: |
-        black --check src
-        ruff check src
+        ruff check --show-fixes
+        ruff format --check
         mypy src
 
     - name: Run tests
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.11.13
+    hooks:
+      - id: ruff-check
+        args: [--fix, --show-fixes]
+      - id: ruff-format
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,7 +37,6 @@ dependencies = [
 test = [
     "mypy",
     "ruff",
-    "black",
     "pytest",
     "pytest-asyncio",
     "pytest-rerunfailures",
@@ -71,3 +70,4 @@ Homepage = "https://github.com/davidbrochart/akernel"
 
 [tool.ruff]
 line-length = 100
+exclude = ["examples"]
diff --git a/src/akernel/cache.py b/src/akernel/cache.py
@@ -10,9 +10,7 @@
 
 def cache(cache_dir: str | None):
     if not cache_dir:
-        cache_dir = os.path.join(
-            sys.prefix, "share", "jupyter", "kernels", "akernel", "cache"
-        )
+        cache_dir = os.path.join(sys.prefix, "share", "jupyter", "kernels", "akernel", "cache")
 
     l4 = File(cache_dir)
     l3 = Func(zlib.compress, zlib.decompress, l4)
diff --git a/src/akernel/code.py b/src/akernel/code.py
@@ -92,9 +92,7 @@ def get_return_body(val):
 
 
 class Transform:
-    def __init__(
-        self, code: str, task_i: int | None = None, react: bool = False
-    ) -> None:
+    def __init__(self, code: str, task_i: int | None = None, react: bool = False) -> None:
         self.gtree = gast.parse(code)
         self.task_i = task_i
         self.react = react
@@ -120,9 +118,7 @@ def get_async_ast(self) -> gast.Module:
             new_body += self.gtree.body + body_globals_update_locals + last_statement
         else:
             new_body += self.gtree.body + body_globals_update_locals
-        name = (
-            "__async_cell__" if self.task_i is None else f"__async_cell{self.task_i}__"
-        )
+        name = "__async_cell__" if self.task_i is None else f"__async_cell{self.task_i}__"
         body = [
             gast.AsyncFunctionDef(
                 name=name,
@@ -169,16 +165,12 @@ def make_react(self):
                         ):
                             # RHS
                             rhs_calls = [
-                                n
-                                for n in gast.walk(statement.value)
-                                if isinstance(n, gast.Call)
+                                n for n in gast.walk(statement.value) if isinstance(n, gast.Call)
                             ]
                             for n in rhs_calls:
                                 ipyx_name = gast.Name(id="ipyx", ctx=gast.Load())
                                 n.func = gast.Call(
-                                    func=gast.Attribute(
-                                        value=ipyx_name, attr="F", ctx=gast.Load()
-                                    ),
+                                    func=gast.Attribute(value=ipyx_name, attr="F", ctx=gast.Load()),
                                     args=[n.func],
                                     keywords=[],
                                 )
@@ -243,9 +235,7 @@ def visit_Name(self, node):
     def visit_Assign(self, node):
         ctx, g = self.context[-1]
         if ctx == "global":
-            self.outputs += [
-                target.id for target in node.targets if isinstance(target, gast.Name)
-            ]
+            self.outputs += [target.id for target in node.targets if isinstance(target, gast.Name)]
         self.generic_visit(node)
 
     def visit_AugAssign(self, node):
diff --git a/src/akernel/comm/comm.py b/src/akernel/comm/comm.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from typing import Dict, List, Any, Callable
+from typing import Any, Callable
 
-import comm  # type: ignore
+import comm
 
 from ..message import send_message, create_message
 
@@ -11,9 +11,9 @@ class Comm(comm.base_comm.BaseComm):
     _msg_callback: Callable | None
     comm_id: str
     topic: bytes
-    parent_header: Dict[str, Any]
+    parent_header: dict[str, Any]
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         from akernel.kernel import KERNEL, PARENT_VAR, Kernel
 
         self.kernel: Kernel = KERNEL
@@ -23,10 +23,10 @@ def __init__(self, **kwargs):
     def publish_msg(
         self,
         msg_type: str,
-        data: Dict[str, Any],
-        metadata: Dict[str, Any],
-        buffers: List[bytes],
-        **keys,
+        data: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+        buffers: list[bytes] | None = None,
+        **keys: Any,
     ) -> None:
         msg = create_message(
             msg_type,
@@ -42,7 +42,7 @@ def publish_msg(
             buffers=buffers,
         )
 
-    def handle_msg(self, msg: Dict[str, Any]) -> None:
+    def handle_msg(self, msg: dict[str, Any]) -> None:
         if self._msg_callback:
             self.kernel.execution_state = "busy"
             msg2 = create_message(
diff --git a/src/akernel/comm/manager.py b/src/akernel/comm/manager.py
@@ -1,21 +1,22 @@
-from typing import Dict, Callable
+from typing import Dict, Callable, cast
 
-import comm  # type: ignore
+import comm
 
 from .comm import Comm
 
 
 class CommManager(comm.CommManager):
-    comms: Dict[str, Comm]
+    comms: dict[str, comm.base_comm.BaseComm]
     targets: Dict[str, Callable]
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         from akernel.kernel import KERNEL, Kernel
 
         self.kernel: Kernel = KERNEL
 
-    def register_comm(self, comm: Comm) -> str:
+    def register_comm(self, comm: comm.base_comm.BaseComm) -> str:
+        comm = cast(Comm, comm)
         comm_id = comm.comm_id
         comm.kernel = self.kernel
         self.comms[comm_id] = comm
diff --git a/src/akernel/execution.py b/src/akernel/execution.py
@@ -43,8 +43,7 @@ def pre_execute(
             f"{Style.RESET_ALL}:",
             f"{Fore.RED}{exception.text.rstrip()}{Style.RESET_ALL}",
             (exception.offset - 1) * " " + "^",
-            f"{Fore.RED}{type(exception).__name__}{Style.RESET_ALL}: "
-            f"{exception.args[0]}",
+            f"{Fore.RED}{type(exception).__name__}{Style.RESET_ALL}: {exception.args[0]}",
         ]
     else:
         if cache is not None:
diff --git a/src/akernel/kernel.py b/src/akernel/kernel.py
@@ -102,9 +102,7 @@ def get_comm_manager():
         self.iopub_channel = connect_channel("iopub", self.connection_cfg)
         self.control_channel = connect_channel("control", self.connection_cfg)
         self.stdin_channel = connect_channel("stdin", self.connection_cfg)
-        msg = self.create_message(
-            "status", content={"execution_state": self.execution_state}
-        )
+        msg = self.create_message("status", content={"execution_state": self.execution_state})
         send_message(msg, self.iopub_channel, self.key)
         self.execution_state = "idle"
         self.stop = asyncio.Event()
@@ -166,8 +164,7 @@ def init_kernel(self, namespace):
         self.locals[namespace] = {}
         if self.react_kernel:
             code = (
-                "import ipyx, ipywidgets;"
-                "globals().update({'ipyx': ipyx, 'ipywidgets': ipywidgets})"
+                "import ipyx, ipywidgets;globals().update({'ipyx': ipyx, 'ipywidgets': ipywidgets})"
             )
             exec(code, self.globals[namespace], self.locals[namespace])
 
@@ -314,10 +311,7 @@ async def listen_shell(self) -> None:
                         parent_header=parent_header,
                         content={
                             "status": "ok",
-                            "comms": {
-                                comm_id: {"target_name": target_name}
-                                for comm_id in comms
-                            },
+                            "comms": {comm_id: {"target_name": target_name} for comm_id in comms},
                         },
                     )
                     send_message(msg2, self.shell_channel, self.key, idents[0])
@@ -329,7 +323,7 @@ async def listen_shell(self) -> None:
                 )
                 send_message(msg2, self.iopub_channel, self.key)
             elif msg_type == "comm_msg":
-                self.comm_manager.comm_msg(None, None, msg)
+                self.comm_manager.comm_msg(None, None, msg)  # type: ignore[arg-type]
 
     async def listen_control(self) -> None:
         while True:
diff --git a/src/akernel/kernelspec.py b/src/akernel/kernelspec.py
@@ -5,9 +5,7 @@
 import json
 
 
-def write_kernelspec(
-    dir_name: str, mode: str, display_name: str, cache_dir: str | None
-) -> None:
+def write_kernelspec(dir_name: str, mode: str, display_name: str, cache_dir: str | None) -> None:
     argv = ["akernel", "launch"]
     if mode:
         argv.append(mode)
diff --git a/src/akernel/message.py b/src/akernel/message.py
@@ -4,7 +4,7 @@
 import hmac
 import hashlib
 from datetime import datetime, timezone
-from typing import List, Dict, Tuple, Any, cast
+from typing import List, Tuple, Any, cast
 
 from zmq.utils import jsonapi
 from zmq.asyncio import Socket
@@ -17,8 +17,8 @@
 DELIM = b"<IDS|MSG>"
 
 
-def date_to_str(obj: Dict[str, Any]):
-    if obj is not None and "date" in obj and type(obj["date"]) != str:
+def date_to_str(obj: dict[str, Any]):
+    if obj is not None and "date" in obj and not isinstance(obj["date"], str):
         obj["date"] = obj["date"].isoformat().replace("+00:00", "Z")
     return obj
 
@@ -32,9 +32,7 @@ def feed_identities(msg_list: List[bytes]) -> Tuple[List[bytes], List[bytes]]:
     return msg_list[:idx], msg_list[idx + 1 :]  # noqa
 
 
-def create_message_header(
-    msg_type: str, session_id: str, msg_cnt: int
-) -> Dict[str, Any]:
+def create_message_header(msg_type: str, session_id: str, msg_cnt: int) -> dict[str, Any]:
     if not session_id:
         session_id = msg_id = uuid.uuid4().hex
     else:
@@ -52,12 +50,12 @@ def create_message_header(
 
 def create_message(
     msg_type: str,
-    content: Dict = {},
-    metadata: Dict = {},
-    parent_header: Dict[str, Any] = {},
+    content: dict = {},
+    metadata: dict[str, Any] | None = None,
+    parent_header: dict[str, Any] = {},
     session_id: str = "",
     msg_cnt: int = 0,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     if parent_header:
         session_id = parent_header["session"]
     header = create_message_header(msg_type, session_id, msg_cnt)
@@ -72,7 +70,7 @@ def create_message(
     return msg
 
 
-def serialize(msg: Dict[str, Any], key: str, address: bytes = b"") -> List[bytes]:
+def serialize(msg: dict[str, Any], key: str, address: bytes = b"") -> List[bytes]:
     message = [
         pack(date_to_str(msg["header"])),
         pack(date_to_str(msg["parent_header"])),
@@ -88,7 +86,7 @@ def serialize(msg: Dict[str, Any], key: str, address: bytes = b"") -> List[bytes
 
 async def receive_message(
     sock: Socket, timeout: float = float("inf")
-) -> Tuple[List[bytes], Dict[str, Any]] | None:
+) -> Tuple[List[bytes], dict[str, Any]] | None:
     timeout *= 1000  # in ms
     ready = await sock.poll(timeout)
     if ready:
@@ -99,7 +97,7 @@ async def receive_message(
 
 
 def send_message(
-    msg: Dict[str, Any],
+    msg: dict[str, Any],
     sock: Socket,
     key: str,
     address: bytes = b"",
@@ -117,12 +115,12 @@ def send_message(
     sock.send_multipart(to_send, copy=True)
 
 
-def pack(obj: Dict[str, Any]) -> bytes:
+def pack(obj: dict[str, Any]) -> bytes:
     return jsonapi.dumps(obj)
 
 
-def unpack(s: bytes) -> Dict[str, Any]:
-    return cast(Dict[str, Any], jsonapi.loads(s))
+def unpack(s: bytes) -> dict[str, Any]:
+    return cast(dict[str, Any], jsonapi.loads(s))
 
 
 def sign(msg_list: List[bytes], key: str) -> bytes:
@@ -133,14 +131,14 @@ def sign(msg_list: List[bytes], key: str) -> bytes:
     return h.hexdigest().encode()
 
 
-def str_to_date(obj: Dict[str, Any]) -> Dict[str, Any]:
+def str_to_date(obj: dict[str, Any]) -> dict[str, Any]:
     if "date" in obj:
         obj["date"] = dateutil_parse(obj["date"])
     return obj
 
 
-def deserialize(msg_list: List[bytes]) -> Dict[str, Any]:
-    message: Dict[str, Any] = {}
+def deserialize(msg_list: List[bytes]) -> dict[str, Any]:
+    message: dict[str, Any] = {}
     header = unpack(msg_list[1])
     message["header"] = str_to_date(header)
     message["msg_id"] = header["msg_id"]
@@ -152,5 +150,5 @@ def deserialize(msg_list: List[bytes]) -> Dict[str, Any]:
     return message
 
 
-async def check_message(sock: Socket) -> int:
+async def check_message(sock: Socket) -> list[tuple[Any, int]]:
     return await sock.poll(0)
diff --git a/src/akernel/traceback.py b/src/akernel/traceback.py
@@ -43,8 +43,5 @@ def get_traceback(code: str, exception, execution_count: int = 0):
         ]
         trace.append(code.splitlines()[frame.f_lineno - 1])
         traceback += trace
-    traceback += [
-        f"{Fore.RED}{type(exception).__name__}{Style.RESET_ALL}: "
-        f"{exception.args[0]}"
-    ]
+    traceback += [f"{Fore.RED}{type(exception).__name__}{Style.RESET_ALL}: {exception.args[0]}"]
     return traceback
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -17,9 +17,7 @@ def all_modes(request):
         mode = "-".join(modes)
         kernel_name += f"-{mode}"
     if mode == "cache":
-        cache_dir = os.path.join(
-            sys.prefix, "share", "jupyter", "kernels", "akernel", "cache"
-        )
+        cache_dir = os.path.join(sys.prefix, "share", "jupyter", "kernels", "akernel", "cache")
         shutil.rmtree(cache_dir, ignore_errors=True)
     display_name = f"Python 3 ({kernel_name})"
     write_kernelspec(kernel_name, mode, display_name, None)
diff --git a/tests/test_kernel.py b/tests/test_kernel.py
diff --git a/tests/test_react_code.py b/tests/test_react_code.py
