Update typing for nbgraderformat

Author: jhamrick

nbgrader/nbgraderformat/common.py

@@ -32,7 +32,7 @@ def __init__(self) -> None:
         with open(os.path.join(root, "v{:d}.json".format(self.schema_version)), "r") as fh:
             self.schema = json.loads(fh.read())
 
-    def _remove_extra_keys(self, cell):
+    def _remove_extra_keys(self, cell: NotebookNode) -> None:
         meta = cell.metadata['nbgrader']
         allowed = set(self.schema["properties"].keys())
         keys = set(meta.keys()) - allowed
@@ -41,12 +41,12 @@ def _remove_extra_keys(self, cell):
             for key in keys:
                 del meta[key]
 
-    def upgrade_notebook_metadata(self, nb):
+    def upgrade_notebook_metadata(self, nb: NotebookNode) -> NotebookNode:
         for cell in nb.cells:
             self.upgrade_cell_metadata(cell)
         return nb
 
-    def upgrade_cell_metadata(self, cell):  # pragma: no cover
+    def upgrade_cell_metadata(self, cell: NotebookNode) -> NotebookNode:  # pragma: no cover
         raise NotImplementedError("this method must be implemented by subclasses")
 
     def validate_cell(self, cell: NotebookNode) -> None:

nbgrader/nbgraderformat/v1.py

@@ -1,6 +1,8 @@
 from nbformat import read as _read, reads as _reads
 from nbformat import write as _write, writes as _writes
+from nbformat.notebooknode import NotebookNode
 from .common import BaseMetadataValidator, ValidationError
+import typing
 
 
 class MetadataValidatorV1(BaseMetadataValidator):
@@ -10,7 +12,7 @@ class MetadataValidatorV1(BaseMetadataValidator):
     def __init__(self) -> None:
         super().__init__()
 
-    def _upgrade_v0_to_v1(self, cell):
+    def _upgrade_v0_to_v1(self, cell: NotebookNode) -> NotebookNode:
         meta = cell.metadata['nbgrader']
 
         if 'grade' not in meta:
@@ -43,7 +45,7 @@ def _upgrade_v0_to_v1(self, cell):
 
         return cell
 
-    def upgrade_cell_metadata(self, cell):
+    def upgrade_cell_metadata(self, cell: NotebookNode) -> NotebookNode:
         if 'nbgrader' not in cell.metadata:
             return cell
 
@@ -59,7 +61,7 @@ def upgrade_cell_metadata(self, cell):
         self._remove_extra_keys(cell)
         return cell
 
-    def validate_cell(self, cell):
+    def validate_cell(self, cell: NotebookNode) -> None:
         super(MetadataValidatorV1, self).validate_cell(cell)
 
         if 'nbgrader' not in cell.metadata:
@@ -92,7 +94,7 @@ def validate_cell(self, cell):
             raise ValidationError(
                 "Markdown solution cell is not marked as a grade cell: {}".format(cell.source))
 
-    def validate_nb(self, nb):
+    def validate_nb(self, nb: NotebookNode) -> None:
         super(MetadataValidatorV1, self).validate_nb(nb)
 
         ids = set([])
@@ -114,23 +116,23 @@ def validate_nb(self, nb):
             ids.add(grade_id)
 
 
-def read_v1(source, as_version, **kwargs):
+def read_v1(source: typing.io.TextIO, as_version: int, **kwargs: typing.Any) -> NotebookNode:
     nb = _read(source, as_version, **kwargs)
     MetadataValidatorV1().validate_nb(nb)
     return nb
 
 
-def write_v1(nb, fp, **kwargs):
+def write_v1(nb: NotebookNode, fp: typing.io.TextIO, **kwargs: typing.Any) -> None:
     MetadataValidatorV1().validate_nb(nb)
-    return _write(nb, fp, **kwargs)
+    _write(nb, fp, **kwargs)
 
 
-def reads_v1(source, as_version, **kwargs):
+def reads_v1(source: str, as_version: int, **kwargs: typing.Any) -> NotebookNode:
     nb = _reads(source, as_version, **kwargs)
     MetadataValidatorV1().validate_nb(nb)
     return nb
 
 
-def writes_v1(nb, **kwargs):
+def writes_v1(nb: NotebookNode, **kwargs: typing.Any) -> None:
     MetadataValidatorV1().validate_nb(nb)
-    return _writes(nb, **kwargs)
+    _writes(nb, **kwargs)

nbgrader/nbgraderformat/v2.py

@@ -1,5 +1,7 @@
 from nbformat import read as _read, reads as _reads
 from nbformat import write as _write, writes as _writes
+from nbformat.notebooknode import NotebookNode
+import typing
 from .v1 import MetadataValidatorV1
 from .common import BaseMetadataValidator, ValidationError
 
@@ -12,7 +14,7 @@ def __init__(self) -> None:
         super().__init__()
         self.v1 = MetadataValidatorV1()
 
-    def _upgrade_v1_to_v2(self, cell):
+    def _upgrade_v1_to_v2(self, cell: NotebookNode) -> NotebookNode:
         meta = cell.metadata['nbgrader']
 
         # only add cell type if the checksum has also already been set
@@ -24,7 +26,7 @@ def _upgrade_v1_to_v2(self, cell):
 
         return cell
 
-    def upgrade_cell_metadata(self, cell):
+    def upgrade_cell_metadata(self, cell: NotebookNode) -> NotebookNode:
         if 'nbgrader' not in cell.metadata:
             return cell
 
@@ -42,7 +44,7 @@ def upgrade_cell_metadata(self, cell):
         self._remove_extra_keys(cell)
         return cell
 
-    def validate_cell(self, cell):
+    def validate_cell(self, cell: NotebookNode) -> None:
         super(MetadataValidatorV2, self).validate_cell(cell)
 
         if 'nbgrader' not in cell.metadata:
@@ -81,7 +83,7 @@ def validate_cell(self, cell):
             raise ValidationError(
                 "Markdown solution cell is not marked as a grade cell: {}".format(cell.source))
 
-    def validate_nb(self, nb):
+    def validate_nb(self, nb: NotebookNode) -> None:
         super(MetadataValidatorV2, self).validate_nb(nb)
 
         ids = set([])
@@ -103,23 +105,23 @@ def validate_nb(self, nb):
             ids.add(grade_id)
 
 
-def read_v2(source, as_version, **kwargs):
+def read_v2(source: typing.io.TextIO, as_version: int, **kwargs: typing.Any) -> NotebookNode:
     nb = _read(source, as_version, **kwargs)
     MetadataValidatorV2().validate_nb(nb)
     return nb
 
 
-def write_v2(nb, fp, **kwargs):
+def write_v2(nb: NotebookNode, fp: typing.io.TextIO, **kwargs: typing.Any) -> None:
     MetadataValidatorV2().validate_nb(nb)
-    return _write(nb, fp, **kwargs)
+    _write(nb, fp, **kwargs)
 
 
-def reads_v2(source, as_version, **kwargs):
+def reads_v2(source: str, as_version: int, **kwargs: typing.Any) -> NotebookNode:
     nb = _reads(source, as_version, **kwargs)
     MetadataValidatorV2().validate_nb(nb)
     return nb
 
 
-def writes_v2(nb, **kwargs):
+def writes_v2(nb: NotebookNode, **kwargs: typing.Any) -> None:
     MetadataValidatorV2().validate_nb(nb)
-    return _writes(nb, **kwargs)
+    _writes(nb, **kwargs)

nbgrader/nbgraderformat/v3.py

@@ -1,5 +1,7 @@
 from nbformat import read as _read, reads as _reads
 from nbformat import write as _write, writes as _writes
+from nbformat.notebooknode import NotebookNode
+import typing
 from .v1 import MetadataValidatorV1
 from .v2 import MetadataValidatorV2
 from .common import BaseMetadataValidator, ValidationError
@@ -15,13 +17,13 @@ def __init__(self) -> None:
         self.v1 = MetadataValidatorV1()
         self.v2 = MetadataValidatorV2()
 
-    def _upgrade_v2_to_v3(self, cell):
+    def _upgrade_v2_to_v3(self, cell: NotebookNode) -> NotebookNode:
         meta = cell.metadata['nbgrader']
         meta['schema_version'] = self.schema_version
 
         return cell
 
-    def upgrade_cell_metadata(self, cell):
+    def upgrade_cell_metadata(self, cell: NotebookNode) -> NotebookNode:
         if 'nbgrader' not in cell.metadata:
             return cell
 
@@ -109,23 +111,23 @@ def validate_nb(self, nb: NotebookNode) -> None:
             ids.add(grade_id)
 
 
-def read_v3(source, as_version, **kwargs):
+def read_v3(source: typing.io.TextIO, as_version: int, **kwargs: typing.Any) -> NotebookNode:
     nb = _read(source, as_version, **kwargs)
     MetadataValidatorV3().validate_nb(nb)
     return nb
 
 
-def write_v3(nb, fp, **kwargs):
+def write_v3(nb: NotebookNode, fp: typing.io.TextIO, **kwargs: typing.Any) -> None:
     MetadataValidatorV3().validate_nb(nb)
-    return _write(nb, fp, **kwargs)
+    _write(nb, fp, **kwargs)
 
 
-def reads_v3(source, as_version, **kwargs):
+def reads_v3(source: str, as_version: int, **kwargs: typing.Any) -> NotebookNode:
     nb = _reads(source, as_version, **kwargs)
     MetadataValidatorV3().validate_nb(nb)
     return nb
 
 
-def writes_v3(nb, **kwargs):
+def writes_v3(nb: NotebookNode, **kwargs: typing.Any) -> None:
     MetadataValidatorV3().validate_nb(nb)
-    return _writes(nb, **kwargs)
+    _writes(nb, **kwargs)