Add some type annotations (#155)

nteract · Aug 1, 2024 · 6f220ed · 6f220ed
1 parent 54fe1d3
commit 6f220ed
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/testbook/testbook.py b/testbook/testbook.py
@@ -47,7 +47,7 @@ def _prepare(self):
         elif self.execute not in [None, False]:
             self.client.execute_cell(self.execute)
 
-    def __enter__(self):
+    def __enter__(self) -> TestbookNotebookClient:
         with self.client.setup_kernel(cleanup_kc=False):
             self._prepare()
             return self.client

diff --git a/testbook/testbooknode.py b/testbook/testbooknode.py
@@ -10,7 +10,7 @@ def __init__(self, *args, **kw):
         super().__init__(*args, **kw)
 
     @property
-    def output_text(self):
+    def output_text(self) -> str:
         text = ''
         for output in self['outputs']:
             if 'text' in output: