improve typing

Author: FynnBe

.pre-commit-config.yaml

@@ -4,10 +4,11 @@ repos:
     hooks:
     - id: black
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.961
+    rev: v0.991
     hooks:
       - id: mypy
         additional_dependencies: [types-requests]
+        args: [--install-types, --non-interactive, --explicit-package-bases, --check-untyped-defs]
   - repo: local
     hooks:
     - id: generate rdf docs

bioimageio/spec/__main__.py

@@ -1,11 +1,14 @@
 import sys
 from pathlib import Path
 from pprint import pprint
+from typing import Any, Callable, Dict, Optional, Union
 
 import typer
 
 from bioimageio.spec import __version__, collection, commands, model, rdf
+from spec.shared.raw_nodes import URI
 
+enrich_partial_rdf_with_imjoy_plugin: Optional[Callable[[Dict[str, Any], Union[URI, Path]], Dict[str, Any]]]
 try:
     from bioimageio.spec.partner.utils import enrich_partial_rdf_with_imjoy_plugin
 except ImportError:
@@ -81,6 +84,7 @@ def validate_partner_collection(
         ),
         verbose: bool = typer.Option(False, help="show traceback of unexpected (no ValidationError) exceptions"),
     ):
+        assert enrich_partial_rdf_with_imjoy_plugin is not None
         summary = commands.validate(
             rdf_source, update_format, update_format_inner, enrich_partial_rdf=enrich_partial_rdf_with_imjoy_plugin
         )
@@ -101,10 +105,12 @@ def validate_partner_collection(
 
         sys.exit(ret_code)
 
+    cmd_doc = commands.validate.__doc__
+    assert cmd_doc is not None
     validate_partner_collection.__doc__ = (
         "A special version of the bioimageio validate command that enriches the RDFs defined in collections by parsing any "
         "associated imjoy plugins. This is implemented using the 'enrich_partial_rdf' of the regular validate command:\n"
-        + commands.validate.__doc__
+        + cmd_doc
     )
 
 

bioimageio/spec/collection/v0_2/schema.py

@@ -1,18 +1,21 @@
-from marshmallow import missing, validates
+from types import ModuleType
+from typing import ClassVar, List, Union
+
+from marshmallow import INCLUDE, missing, validates
 
 from bioimageio.spec.rdf.v0_2.schema import RDF
 from bioimageio.spec.shared import fields
 from bioimageio.spec.shared.schema import SharedBioImageIOSchema, WithUnknown
 from . import raw_nodes
 
 try:
-    from typing import List, Union, get_args
+    from typing import get_args
 except ImportError:
     from typing_extensions import get_args  # type: ignore
 
 
 class _BioImageIOSchema(SharedBioImageIOSchema):
-    raw_nodes = raw_nodes
+    raw_nodes: ClassVar[ModuleType] = raw_nodes
 
 
 class CollectionEntry(_BioImageIOSchema, WithUnknown):
@@ -21,6 +24,9 @@ class CollectionEntry(_BioImageIOSchema, WithUnknown):
 
 
 class Collection(_BioImageIOSchema, WithUnknown, RDF):
+    class Meta:
+        unknown = INCLUDE
+
     bioimageio_description = f"""# BioImage.IO Collection Resource Description File Specification {get_args(raw_nodes.FormatVersion)[-1]}
 This specification defines the fields used in a BioImage.IO-compliant resource description file (`RDF`) for describing collections of other resources.
 These fields are typically stored in a YAML file which we call Collection Resource Description File or `collection RDF`.

bioimageio/spec/commands.py

@@ -32,7 +32,7 @@ def update_format(
 def validate(
     rdf_source: Union[RawResourceDescription, dict, os.PathLike, IO, str, bytes],
     update_format: bool = False,
-    update_format_inner: bool = None,
+    update_format_inner: Optional[bool] = None,
     verbose: bool = "deprecated",  # type: ignore
     enrich_partial_rdf: Callable[[dict, Union[URI, Path]], dict] = default_enrich_partial_rdf,
 ) -> ValidationSummary:
@@ -148,7 +148,7 @@ def validate(
 def update_rdf(
     source: Union[RawResourceDescription, dict, os.PathLike, IO, str, bytes],
     update: Union[RawResourceDescription, dict, os.PathLike, IO, str, bytes],
-    output: Union[dict, os.PathLike] = None,
+    output: Union[None, dict, os.PathLike] = None,
     validate_output: bool = True,
 ) -> Union[dict, Path, RawResourceDescription]:
     """
@@ -215,6 +215,7 @@ def update_rdf(
                     f"Failed to convert paths in updated rdf to relative paths with root {output}; error: {e}"
                 )
                 warnings.warn(f"updated rdf at {output} contains absolute paths and is thus invalid!")
+                assert isinstance(out_data, RawResourceDescription)
                 out_data = serialize_raw_resource_description_to_dict(out_data, convert_absolute_paths=False)
 
         yaml.dump(out_data, output)

bioimageio/spec/dataset/v0_2/schema.py

@@ -1,3 +1,6 @@
+from types import ModuleType
+from typing import ClassVar
+
 from bioimageio.spec.rdf.v0_2.schema import RDF
 from bioimageio.spec.shared.schema import SharedBioImageIOSchema
 from . import raw_nodes
@@ -9,7 +12,7 @@
 
 
 class _BioImageIOSchema(SharedBioImageIOSchema):
-    raw_nodes = raw_nodes
+    raw_nodes: ClassVar[ModuleType] = raw_nodes
 
 
 class Dataset(_BioImageIOSchema, RDF):

bioimageio/spec/io_.py

@@ -81,7 +81,7 @@ def extract_resource_package(
     source: Union[os.PathLike, IO, str, bytes, raw_nodes.URI]
 ) -> Tuple[dict, str, pathlib.Path]:
     """extract a zip source to BIOIMAGEIO_CACHE_PATH"""
-    source, source_name, root = resolve_rdf_source(source)
+    src, source_name, root = resolve_rdf_source(source)
     if isinstance(root, bytes):
         raise NotImplementedError("package source was bytes")
 
@@ -123,7 +123,7 @@ def extract_resource_package(
             warnings.warn(f"Could not remove download {download} due to {e}")
 
     assert isinstance(package_path, pathlib.Path)
-    return source, source_name, package_path
+    return src, source_name, package_path
 
 
 def load_raw_resource_description(
@@ -292,7 +292,7 @@ def get_resource_package_content_wo_rdf(
     raw_rd: Union[GenericRawRD, raw_nodes.URI, str, pathlib.Path],
     *,
     weights_priority_order: Optional[Sequence[str]] = None,  # model only
-) -> Tuple[GenericRawNode, Dict[str, Union[pathlib.PurePath, raw_nodes.URI]]]:
+) -> Tuple[raw_nodes.ResourceDescription, Dict[str, Union[pathlib.PurePath, raw_nodes.URI]]]:
     """
     Args:
         raw_rd: raw resource description
@@ -305,25 +305,27 @@ def get_resource_package_content_wo_rdf(
         keyed by file names.
         Important note: the serialized rdf.yaml is not included.
     """
-    if not isinstance(raw_rd, raw_nodes.ResourceDescription):
-        raw_rd = load_raw_resource_description(raw_rd)
+    if isinstance(raw_rd, raw_nodes.ResourceDescription):
+        r_rd = raw_rd
+    else:
+        r_rd = load_raw_resource_description(raw_rd)
 
-    sub_spec = _get_spec_submodule(raw_rd.type, raw_rd.format_version)
-    if raw_rd.type == "model":
+    sub_spec = _get_spec_submodule(r_rd.type, r_rd.format_version)
+    if r_rd.type == "model":
         filter_kwargs = dict(weights_priority_order=weights_priority_order)
     else:
         filter_kwargs = {}
 
-    raw_rd = sub_spec.utils.filter_resource_description(raw_rd, **filter_kwargs)
+    r_rd = sub_spec.utils.filter_resource_description(r_rd, **filter_kwargs)
 
-    content: Dict[str, Union[pathlib.PurePath, raw_nodes.URI, str]] = {}
-    raw_rd = RawNodePackageTransformer(content, raw_rd.root_path).transform(raw_rd)
+    content: Dict[str, Union[pathlib.PurePath, raw_nodes.URI]] = {}
+    r_rd = RawNodePackageTransformer(content, r_rd.root_path).transform(r_rd)
     assert "rdf.yaml" not in content
-    return raw_rd, content
+    return r_rd, content
 
 
 def get_resource_package_content(
-    raw_rd: Union[GenericRawNode, raw_nodes.URI, str, pathlib.Path],
+    raw_rd: Union[raw_nodes.ResourceDescription, raw_nodes.URI, str, pathlib.Path],
     *,
     weights_priority_order: Optional[Sequence[str]] = None,  # model only
 ) -> Dict[str, Union[str, pathlib.PurePath, raw_nodes.URI]]:
@@ -343,7 +345,5 @@ def get_resource_package_content(
             "without yaml"
         )
 
-    content: Dict[str, Union[str, pathlib.PurePath, raw_nodes.URI]]
-    raw_rd, content = get_resource_package_content_wo_rdf(raw_rd, weights_priority_order=weights_priority_order)
-    content["rdf.yaml"] = serialize_raw_resource_description(raw_rd)
-    return content
+    r_rd, content = get_resource_package_content_wo_rdf(raw_rd, weights_priority_order=weights_priority_order)
+    return {**content, **{"rdf.yaml": serialize_raw_resource_description(r_rd)}}

bioimageio/spec/model/v0_3/converters.py

@@ -1,6 +1,6 @@
 import copy
 import pathlib
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 from marshmallow import Schema
 
@@ -62,7 +62,7 @@ class DocSchema(Schema):
                 assert isinstance(orig_doc, str)
                 if orig_doc.startswith("http"):
                     if orig_doc.endswith(".md"):
-                        doc = raw_nodes.URI(orig_doc)
+                        doc: Union[raw_nodes.URI, str, pathlib.Path] = raw_nodes.URI(orig_doc)
                     else:
                         doc = f"Find documentation at {orig_doc}"
                 else:

bioimageio/spec/model/v0_3/schema.py

@@ -1,6 +1,7 @@
 import typing
 import warnings
 from copy import deepcopy
+from types import ModuleType
 
 from marshmallow import RAISE, ValidationError, missing as missing_, post_load, pre_dump, pre_load, validates_schema
 
@@ -21,7 +22,7 @@
 
 
 class _BioImageIOSchema(SharedBioImageIOSchema):
-    raw_nodes = raw_nodes
+    raw_nodes: typing.ClassVar[ModuleType] = raw_nodes
 
 
 class RunMode(_BioImageIOSchema):
@@ -474,7 +475,7 @@ class ModelParent(_BioImageIOSchema):
 
 
 class Model(rdf.schema.RDF):
-    raw_nodes = raw_nodes
+    raw_nodes: typing.ClassVar[ModuleType] = raw_nodes
 
     class Meta:
         unknown = RAISE
@@ -758,7 +759,7 @@ def validate_reference_tensor_names(self, data, **kwargs):
                     raise ValidationError(f"{ref_tensor} not found in inputs")
 
     @validates_schema
-    def weights_entries_match_weights_formats(self, data, **kwargs):
+    def weights_entries_match_weights_formats(self, data, **kwargs) -> None:
         weights: typing.Dict[str, _WeightsEntryBase] = data["weights"]
         for weights_format, weights_entry in weights.items():
             if weights_format in ["keras_hdf5", "tensorflow_js", "tensorflow_saved_model_bundle"]:

bioimageio/spec/model/v0_4/raw_nodes.py

@@ -40,7 +40,9 @@
 
 # reassign to use imported classes
 ImplicitOutputShape = ImplicitOutputShape
+InputTensor = InputTensor
 Maintainer = Maintainer
+OutputTensor = OutputTensor
 ParametrizedInputShape = ParametrizedInputShape
 Postprocessing = Postprocessing
 PostprocessingName = PostprocessingName

bioimageio/spec/model/v0_4/schema.py

@@ -1,5 +1,6 @@
 import typing
 from copy import deepcopy
+from types import ModuleType
 
 import numpy
 from marshmallow import RAISE, ValidationError, missing, pre_load, validates, validates_schema
@@ -23,7 +24,7 @@
 
 
 class _BioImageIOSchema(SharedBioImageIOSchema):
-    raw_nodes = raw_nodes
+    raw_nodes: typing.ClassVar[ModuleType] = raw_nodes
 
 
 class _TensorBase(_BioImageIOSchema):
@@ -184,7 +185,7 @@ def matching_halo_length(self, data, **kwargs):
 
 
 class _WeightsEntryBase(_WeightsEntryBase03):
-    raw_nodes = raw_nodes
+    raw_nodes: typing.ClassVar[ModuleType] = raw_nodes
     dependencies = fields.Dependencies(
         bioimageio_description="Dependency manager and dependency file, specified as `<dependency manager>:<relative "
         "path to file>`. For example: 'conda:./environment.yaml', 'maven:./pom.xml', or 'pip:./requirements.txt'. "
@@ -249,7 +250,7 @@ class Dataset(_Dataset):
 
 
 class TorchscriptWeightsEntry(_WeightsEntryBase):
-    raw_nodes = raw_nodes
+    raw_nodes: typing.ClassVar[ModuleType] = raw_nodes
 
     bioimageio_description = "Torchscript weights format"
     weights_format = fields.String(validate=field_validators.Equal("torchscript"), required=True, load_only=True)
@@ -294,7 +295,7 @@ def id_xor_uri(self, data, **kwargs):
 
 
 class Model(rdf.schema.RDF):
-    raw_nodes = raw_nodes
+    raw_nodes: typing.ClassVar = raw_nodes
 
     class Meta:
         unknown = RAISE
@@ -433,7 +434,7 @@ def no_duplicate_output_tensor_names(self, value: typing.List[raw_nodes.OutputTe
             raise ValidationError("Duplicate output tensor names are not allowed.")
 
     @validates_schema
-    def inputs_and_outputs(self, data, **kwargs):
+    def inputs_and_outputs(self, data, **kwargs) -> None:
         ipts: typing.List[raw_nodes.InputTensor] = data.get("inputs")
         outs: typing.List[raw_nodes.OutputTensor] = data.get("outputs")
         if any(
@@ -603,7 +604,7 @@ def add_weights_format_key_to_weights_entry_value(self, data: dict, many=False,
         return data
 
     @validates_schema
-    def validate_reference_tensor_names(self, data, **kwargs):
+    def validate_reference_tensor_names(self, data, **kwargs) -> None:
         def get_tnames(tname: str):
             return [t.get("name") if isinstance(t, dict) else t.name for t in data.get(tname, [])]
 
@@ -649,7 +650,7 @@ def get_tnames(tname: str):
                     raise ValidationError(f"invalid self reference for preprocessing of tensor {t.name}")
 
     @validates_schema
-    def weights_entries_match_weights_formats(self, data, **kwargs):
+    def weights_entries_match_weights_formats(self, data, **kwargs) -> None:
         weights: typing.Dict[str, WeightsEntry] = data.get("weights", {})
         for weights_format, weights_entry in weights.items():
             if not isinstance(weights_entry, get_args(raw_nodes.WeightsEntry)):