Fix typings and update mypy settings

.github/workflows/main.yml

@@ -54,8 +54,10 @@ jobs:
       run: mypy nbclient
 
     - name: Run the tests
+      shell: bash
       run: |
-        pytest -vv --maxfail=2 --cov=nbclient --cov-report=xml -W always
+        args="-vv --maxfail=2 --cov=nbclient --cov-report=xml -W always"
+        pytest $args || pytest $arg --lf
 
     - name: Upload coverage to Codecov
       uses: codecov/codecov-action@v2

MANIFEST.in

@@ -6,6 +6,7 @@ include *.md
 include tox.ini
 include pyproject.toml
 include .pre-commit-config.yaml
+include nbclient/py.typed
 
 # Code and test files
 recursive-include nbclient *.ipynb

nbclient/cli.py

@@ -4,7 +4,7 @@
 from textwrap import dedent
 
 import nbformat
-from jupyter_core.application import JupyterApp  # type: ignore
+from jupyter_core.application import JupyterApp
 from traitlets import Bool, Integer, List, Unicode, default
 from traitlets.config import catch_config_error
 

nbclient/client.py

@@ -9,7 +9,6 @@
 from queue import Empty
 from textwrap import dedent
 from time import monotonic
-from typing import Optional
 
 from jupyter_client import KernelManager
 from jupyter_client.client import KernelClient
@@ -40,7 +39,7 @@
 from .util import ensure_async, run_hook, run_sync
 
 
-def timestamp(msg: Optional[Dict] = None) -> str:
+def timestamp(msg: t.Optional[Dict] = None) -> str:
     if msg and 'header' in msg:  # The test mocks don't provide a header, so tolerate that
         msg_header = msg['header']
         if 'date' in msg_header and isinstance(msg_header['date'], datetime.datetime):
@@ -78,7 +77,7 @@ class NotebookClient(LoggingConfigurable):
         ),
     ).tag(config=True)
 
-    timeout_func: t.Any = Any(
+    timeout_func: t.Callable[..., t.Optional[int]] = Any(
         default_value=None,
         allow_none=True,
         help=dedent(
@@ -353,8 +352,8 @@ class NotebookClient(LoggingConfigurable):
         ),
     ).tag(config=True)
 
-    @default('kernel_manager_class')
-    def _kernel_manager_class_default(self) -> KernelManager:
+    @default('kernel_manager_class')  # type:ignore[misc]
+    def _kernel_manager_class_default(self) -> t.Type[KernelManager]:
         """Use a dynamic default to avoid importing jupyter_client at startup"""
         from jupyter_client import AsyncKernelManager
 
@@ -403,7 +402,7 @@ def _kernel_manager_class_default(self) -> KernelManager:
         )
     )
 
-    def __init__(self, nb: NotebookNode, km: t.Optional[KernelManager] = None, **kw) -> None:
+    def __init__(self, nb: NotebookNode, km: t.Optional[KernelManager] = None, **kw: t.Any) -> None:
         """Initializes the execution manager.
 
         Parameters
@@ -458,6 +457,7 @@ def create_kernel_manager(self) -> KernelManager:
             self.km = self.kernel_manager_class(config=self.config)
         else:
             self.km = self.kernel_manager_class(kernel_name=self.kernel_name, config=self.config)
+        assert self.km is not None
 
         # If the current kernel manager is still using the default (synchronous) KernelClient class,
         # switch to the async version since that's what NBClient prefers.
@@ -481,13 +481,13 @@ async def _async_cleanup_kernel(self) -> None:
             # Remove any state left over even if we failed to stop the kernel
             await ensure_async(self.km.cleanup_resources())
             if getattr(self, "kc") and self.kc is not None:
-                await ensure_async(self.kc.stop_channels())
+                await ensure_async(self.kc.stop_channels())  # type:ignore
                 self.kc = None
                 self.km = None
 
     _cleanup_kernel = run_sync(_async_cleanup_kernel)
 
-    async def async_start_new_kernel(self, **kwargs) -> None:
+    async def async_start_new_kernel(self, **kwargs: t.Any) -> None:
         """Creates a new kernel.
 
         Parameters
@@ -527,7 +527,7 @@ async def async_start_new_kernel_client(self) -> KernelClient:
         """
         assert self.km is not None
         self.kc = self.km.client()
-        await ensure_async(self.kc.start_channels())
+        await ensure_async(self.kc.start_channels())  # type:ignore[func-returns-value]
         try:
             await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))
         except RuntimeError:
@@ -540,7 +540,7 @@ async def async_start_new_kernel_client(self) -> KernelClient:
     start_new_kernel_client = run_sync(async_start_new_kernel_client)
 
     @contextmanager
-    def setup_kernel(self, **kwargs) -> t.Generator:
+    def setup_kernel(self, **kwargs: t.Any) -> t.Generator:
         """
         Context manager for setting up the kernel to execute a notebook.
 
@@ -567,7 +567,7 @@ def setup_kernel(self, **kwargs) -> t.Generator:
                 self._cleanup_kernel()
 
     @asynccontextmanager
-    async def async_setup_kernel(self, **kwargs) -> t.AsyncGenerator:
+    async def async_setup_kernel(self, **kwargs: t.Any) -> t.AsyncGenerator:
         """
         Context manager for setting up the kernel to execute a notebook.
 
@@ -620,7 +620,7 @@ def on_signal():
             except (NotImplementedError, RuntimeError):
                 pass
 
-    async def async_execute(self, reset_kc: bool = False, **kwargs) -> NotebookNode:
+    async def async_execute(self, reset_kc: bool = False, **kwargs: t.Any) -> NotebookNode:
         """
         Executes each code cell.
 
@@ -727,7 +727,7 @@ async def _async_poll_for_reply(
             new_timeout = float(timeout)
         while True:
             try:
-                msg = await ensure_async(self.kc.shell_channel.get_msg(timeout=new_timeout))
+                msg: t.Dict = await ensure_async(self.kc.shell_channel.get_msg(timeout=new_timeout))
                 if msg['parent_header'].get('msg_id') == msg_id:
                     if self.record_timing:
                         cell['metadata']['execution']['shell.execute_reply'] = timestamp(msg)
@@ -777,7 +777,7 @@ async def _async_poll_kernel_alive(self) -> None:
                 self.task_poll_for_reply.cancel()
                 return
 
-    def _get_timeout(self, cell: t.Optional[NotebookNode]) -> int:
+    def _get_timeout(self, cell: t.Optional[NotebookNode]) -> t.Optional[int]:
         if self.timeout_func is not None and cell is not None:
             timeout = self.timeout_func(cell)
         else:
@@ -818,7 +818,7 @@ async def async_wait_for_reply(
         cummulative_time = 0
         while True:
             try:
-                msg = await ensure_async(
+                msg: t.Dict = await ensure_async(
                     self.kc.shell_channel.get_msg(timeout=self.shell_timeout_interval)
                 )
             except Empty:
@@ -1034,7 +1034,7 @@ def process_message(
 
     def output(
         self, outs: t.List, msg: t.Dict, display_id: str, cell_index: int
-    ) -> t.Optional[t.List]:
+    ) -> t.Optional[NotebookNode]:
 
         msg_type = msg['msg_type']
 
@@ -1167,7 +1167,7 @@ def remove_output_hook(self, msg_id: str, hook: OutputWidget) -> None:
         removed_hook = self.output_hook_stack[msg_id].pop()
         assert removed_hook == hook
 
-    def on_comm_open_jupyter_widget(self, msg: t.Dict):
+    def on_comm_open_jupyter_widget(self, msg: t.Dict) -> t.Optional[t.Any]:
         content = msg['content']
         data = content['data']
         state = data['state']
@@ -1177,11 +1177,15 @@ def on_comm_open_jupyter_widget(self, msg: t.Dict):
             widget_class = module.get(state['_model_name'])
             if widget_class:
                 return widget_class(comm_id, state, self.kc, self)
+        return None
 
 
 def execute(
-    nb: NotebookNode, cwd: t.Optional[str] = None, km: t.Optional[KernelManager] = None, **kwargs
-) -> NotebookClient:
+    nb: NotebookNode,
+    cwd: t.Optional[str] = None,
+    km: t.Optional[KernelManager] = None,
+    **kwargs: t.Any,
+) -> NotebookNode:
     """Execute a notebook's code, updating outputs within the notebook object.
 
     This is a convenient wrapper around NotebookClient. It returns the

nbclient/exceptions.py

@@ -19,7 +19,9 @@ class CellTimeoutError(TimeoutError, CellControlSignal):
     """
 
     @classmethod
-    def error_from_timeout_and_cell(cls, msg: str, timeout: int, cell: NotebookNode):
+    def error_from_timeout_and_cell(
+        cls, msg: str, timeout: int, cell: NotebookNode
+    ) -> "CellTimeoutError":
         if cell and cell.source:
             src_by_lines = cell.source.strip().split("\n")
             src = (
@@ -74,7 +76,7 @@ def __unicode__(self) -> str:
         return self.traceback
 
     @classmethod
-    def from_cell_and_msg(cls, cell: NotebookNode, msg: Dict):
+    def from_cell_and_msg(cls, cell: NotebookNode, msg: Dict) -> "CellExecutionError":
         """Instantiate from a code cell object and a message contents
         (message is either execute_reply or error)
         """

nbclient/output_widget.py

@@ -10,7 +10,7 @@ class OutputWidget:
     """This class mimics a front end output widget"""
 
     def __init__(
-        self, comm_id: str, state: Dict[str, Any], kernel_client: KernelClient, executor
+        self, comm_id: str, state: Dict[str, Any], kernel_client: KernelClient, executor: Any
     ) -> None:
 
         self.comm_id: str = comm_id
@@ -46,7 +46,7 @@ def _publish_msg(
         data: Optional[Dict] = None,
         metadata: Optional[Dict] = None,
         buffers: Optional[List] = None,
-        **keys
+        **keys: Any
     ) -> None:
         """Helper for sending a comm message on IOPub"""
         data = {} if data is None else data

nbclient/tests/test_client.py

@@ -13,12 +13,12 @@
 
 import nbformat
 import pytest
-import xmltodict  # type: ignore
+import xmltodict
 from jupyter_client import KernelManager
 from jupyter_client.kernelspec import KernelSpecManager
-from nbconvert.filters import strip_ansi  # type: ignore
+from nbconvert.filters import strip_ansi
 from nbformat import NotebookNode
-from testpath import modified_env  # type: ignore
+from testpath import modified_env
 from traitlets import TraitError
 
 from .. import NotebookClient, execute
@@ -126,15 +126,15 @@ async def async_run_notebook(filename, opts, resources=None):
     return input_nb, output_nb
 
 
-def prepare_cell_mocks(*messages, reply_msg=None):
+def prepare_cell_mocks(*messages_input, reply_msg=None):
     """
     This function prepares a executor object which has a fake kernel client
     to mock the messages sent over zeromq. The mock kernel client will return
     the messages passed into this wrapper back from ``preproc.kc.iopub_channel.get_msg``
     callbacks. It also appends a kernel idle message to the end of messages.
     """
     parent_id = 'fake_id'
-    messages = list(messages)
+    messages = list(messages_input)
     # Always terminate messages with an idle to exit the loop
     messages.append({'msg_type': 'status', 'content': {'execution_state': 'idle'}})
 
@@ -313,8 +313,10 @@ def test_parallel_notebooks(capfd, tmpdir):
             threading.Thread(target=run_notebook, args=(input_file.format(label=label), opts, res))
             for label in ("A", "B")
         ]
-        [t.start() for t in threads]
-        [t.join(timeout=2) for t in threads]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join(timeout=2)
 
     captured = capfd.readouterr()
     assert filter_messages_on_error_output(captured.err) == ""
@@ -695,7 +697,7 @@ def test_custom_kernel_manager(self):
             self.assertNotEqual(call_count, 0, f'{method} was called')
 
     def test_process_message_wrapper(self):
-        outputs = []
+        outputs: list = []
 
         class WrappedPreProc(NotebookClient):
             def process_message(self, msg, cell, cell_index):

nbclient/tests/test_util.py

@@ -39,10 +39,10 @@ def test_nested_asyncio_with_tornado():
     ioloop = tornado.ioloop.IOLoop.current()
 
     async def some_async_function():
-        future = asyncio.ensure_future(asyncio.sleep(0.1))
+        future: asyncio.Future = asyncio.ensure_future(asyncio.sleep(0.1))
         # this future is a different future after nested-asyncio has patched
         # the asyncio module, check if tornado likes it:
-        ioloop.add_future(future, lambda f: f.result())
+        ioloop.add_future(future, lambda f: f.result())  # type:ignore
         await future
         return 42
 

nbclient/util.py

@@ -13,14 +13,14 @@ def check_ipython() -> None:
     # original from vaex/asyncio.py
     IPython = sys.modules.get('IPython')
     if IPython:
-        version_str = IPython.__version__  # type: ignore
+        version_str = IPython.__version__
         # We get rid of any trailing ".dev"
         version_str = version_str.replace(".dev", "")
 
         IPython_version = tuple(map(int, version_str.split('.')))
         if IPython_version < (7, 0, 0):
             raise RuntimeError(
-                f'You are using IPython {IPython.__version__} '  # type: ignore
+                f'You are using IPython {IPython.__version__} '
                 'while we require 7.0.0+, please update IPython'
             )
 
@@ -29,7 +29,7 @@ def check_patch_tornado() -> None:
     """If tornado is imported, add the patched asyncio.Future to its tuple of acceptable Futures"""
     # original from vaex/asyncio.py
     if 'tornado' in sys.modules:
-        import tornado.concurrent  # type: ignore
+        import tornado.concurrent
 
         if asyncio.Future not in tornado.concurrent.FUTURES:
             tornado.concurrent.FUTURES = tornado.concurrent.FUTURES + (  # type: ignore
@@ -104,7 +104,7 @@ async def ensure_async(obj: Union[Awaitable, Any]) -> Any:
     return obj
 
 
-async def run_hook(hook: Optional[Callable], **kwargs) -> None:
+async def run_hook(hook: Optional[Callable], **kwargs: Any) -> None:
     if hook is None:
         return
     res = hook(**kwargs)

pyproject.toml

@@ -34,14 +34,32 @@ known_first_party = ["nbclient"]
 
 [tool.mypy]
 python_version = 3.9
+check_untyped_defs = true
+disallow_any_generics = false  # todo
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+no_implicit_optional = true
+no_implicit_reexport = false  # todo
+pretty = true
+show_error_context = true
+show_error_codes = true
+strict_equality = true
+strict_optional = true
+warn_unused_configs = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unused_ignores = true
 
 [[tool.mypy.overrides]]
 module = [
+    "async_generator.*",
+    "jupyter_core.*",
     "nbformat.*",
+    "nbconvert.*",
     "nest_asyncio.*",
-    "async_generator.*",
+    "testpath",
     "traitlets.*",
-    "jupyter_client.*",
+    "xmltodict"
 ]
 ignore_missing_imports = true