Improve typing: better type hints, increased type safety, mypy support (

#534)

* Initial mypy configuration

* Fix types: test_envs.py

* Fix types: conftest.py

* Fix types: tests/util

* Fix types: tests/scripts

* Fix types: tests/rewards

* Fix types: tests/policies

* Incorrect decorator in update_stats method form networks.py::BaseNorm

* Fix types: tests/algorithms (adersarial and bc)

* Fix types: tests/algorithms (dagger and pc)

* Fix types: tests/data

* Linting

* Linting

* Fix types: algorithms/preference_comparisons.py

* Fix types: algorithms/mce_irl.py

* Formatting, fixed minor bug

* Clarify why types are ignored

* Started fixing types on algorithms/density.py

* Linting

* Linting (add back type ignore after reformatting)

* Fixed types: imitation/data/types.py

* Fixed types (started): imitation/data/

* Fixed types: imitation/data/buffer.py

* Fixed bug in buffer.py

* Fixed types: imitation/data/rollout.py

* Fixed types: imitation/data/wrappers.py

* Improve makefile to support automatic cache cleaning

* Fixed types: imitation/testing/

* Linting, fixed wrong return type in rewards.predict_processed_all

* Fixed types: imitation/policies/

* Formatting

* Fixed types: imitation/rewards/

* Fixed types: imitation/rewards/

* Fixed types: imitation/scripts/

* Fixed types: imitation/util/ and formatting

* Linting and formatting

* Bug fixes for test errors

* Linting and typing

* Improve typing in algorithms

* Formatting

* Bug fix

* Formatting

* Fixes suggested by Adam.

* Fix mypy version.

* Update TabularPolicy.predict to match base class

* Fix not checking for dones

* Change for loop to dict comprehension

* Remove is_ensemble to clear up type checking errors

* Reduce code duplication and general cleanup

* Fix type annotation of step_dict

* Change List to Sequence

* Fix density.py::DensityAlgorithm._set_demo_from_batch

* Fixed n_steps (OnPolicyAlgorithm)

* Fix errors in tests

* Include some suggestions into rollout.py and preference_comparisons.py

* Formatting

* Fix setter error as per python/mypy#5936

* add reason for assertion.

* Fix style guide violation: https://google.github.io/styleguide/pyguide.html#22-imports

* Update src/imitation/scripts/parallel.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Move kwargs to the end.

* Swap order of expert_policy_type and expert_policy_path validation check

* Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Update tests/rewards/test_reward_fn.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Explicit random state setting and fix corresponding tests (except notebooks, sacred config, scripts)

* Fix notebooks; add script to clean notebooks

* Fix all tests.

* Formattting.

* Additional fixes

* Linting

* Remove automatically generated `_api` docs files too on `make clean`

* Fix docstrings.

* Fix issue with next(iter(iterable))

* Formatting

* Remove whitespace

* Add TODO message to remove type ignore later

* Remove unnecessary assertion.

* Fixed types in density.py set_demonstrations

* Added type ignore to pytype bug

* Fix_get_first_iter_element and add tests

* Bugfix in BC and tests -- masked as previously iterator ran out too early!

* Remove makefile for now

* Added link to SB3 issue for future reference.

* Fix types of train_imitation
Only return "expert_stats" if all trajectories have reward.

* Modify assert in test_bc to reflect correct type

* Add ci/clean_notebooks.py to CI checks

* Improve clean_notebooks.py by allowing checking only mode.

* Add ipynb notebook checks to CI

* Add support for explicit files for notebook cleaning

* Clean notebooks

* Small improvements in util.py

* Replace TransitionKind with TransitionsMinimal

* Delete unused statement in test

* Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Make type ignore specific to pytype

* Linting

* Migrate from RandomState (deprecated) to Generator

* Add backticks to error message

* Create "AnyNorm" alias

* Small fix

* Add additional checks to shapes in _set_demo_from_batch

* Fix RolloutStatsComputer type

* Improved logging/messages in clean_notebooks.py

* Fix issues resulting from merge

* Bug fix

* Bug fix (wasn't really fixed before)

* Fixed docs example of BC

* Fix bugs resulting from merge

* Fix docs (dagger.rst) caught by sphinx CI

* Add mypy to CI

* Continue fixing miscellaneous type errors

* Linting

* Fix issue with normalize_input_layer type

* Add support for checking presence of generic type ignores

* Allow subdirectories in notebook clean

* Add full typing support for TransitionsMinimal as a sequence

* Fix types for density.py

* Misc fixes

* Add support for prefix context manager in logger (from #529)

* Added back accidentally removed code

* Replaced preference comparisons prefix with ctx manager

* Fixed errors

* Bug fixes

* Docstring fixes

* Fix bug in serialize.py

* Fixed codecheck by pointing notebook checks to docs

* Add rng to mce_irl.rst (doctest)

* Add rng to density.rst (doctest)

* Fix remaining rst files

* Increase sample size to reduce flakiness

* Ignore files not passing mypy for now

* Comment in wrong line

* Comment in wrong line

* Move excluded files to argument

* Add quotes to mypy arg call

* Fix CI mypy call

* Fix CI yaml

* Break ignored files up into one line each

* Address PR comments

* Point SB3 to master to include bug fix

* Do not follow imports for ignored files

* Format / fix tests for context manager

* Switch to sb3 1.6.1

* Formatting

* Remove unused import

* Remove unused fixture

* Add coveragerc file

* Add utils test

* Add tests and asserts

* Add test to synthetic gatherer

* Add trajectory unwrap tests

* Formatting

* Remove bracket typo

* Fix .coveragerc instruction

* Improve density algo coverage and bug fixes

* Fix bug in test

* Add pragma no cover updates

* Minor coverage tweaks

* Fix iterator test

* Update ci/check_typeignore.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* DRY clean_notebooks.py

* Minor tweak in check_typeignore.py

* Added all checks to CI

* Move imports to top-level

* Move main to main() function in script

* Minor fixes

* Remove tweak after new SB3 release

* Split main() into helper functions.

* Fix edge case of n=0 in seed generator

* Update src/imitation/scripts/common/rl.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Fix general type ignore in src

* Fix type ignore errors

* Formatting

* Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Update src/imitation/scripts/common/rl.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Update src/imitation/util/util.py

Co-authored-by: Adam Gleave <adam@gleave.me>

* Replace todo with todo+issue link

* Add explicit type ignore arg

* Add excluded files to code_checks.sh

* Unbreak line

* Misc fixes

* Add training to the density algorithm

* Fix no attribute error

* Type ignore to `with raises` test

* Remove unused import

* Check typeignore for all SRC files

* Clean notebooks

* Remove unused import

* Ignore the file itself from typeignore check

* Add exception to docstring

* Fix bad naming for clean_notebooks

Co-authored-by: Adam Gleave <adam@gleave.me>

.circleci/config.yml

@@ -28,15 +28,28 @@ executors:
     resource_class: xlarge
     environment:
       # If you change these, also change ci/code_checks.sh
-      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py
+      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/
       NUM_CPUS: 8
   static-analysis-medium:
     <<: *defaults
     resource_class: medium
     environment:
       # If you change these, also change ci/code_checks.sh
-      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py
+      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/
       NUM_CPUS: 2
+      EXCLUDE_MYPY: | 
+        (?x)(
+          src/imitation/algorithms/preference_comparisons.py$
+          | src/imitation/rewards/reward_nets.py$
+          | src/imitation/util/sacred.py$
+          | src/imitation/algorithms/base.py$
+          | src/imitation/scripts/train_preference_comparisons.py$
+          | src/imitation/rewards/serialize.py$
+          | src/imitation/scripts/common/train.py$
+          | src/imitation/algorithms/mce_irl.py$
+          | src/imitation/algorithms/density.py$
+          | tests/algorithms/test_bc.py$
+        )
 
 commands:
   dependencies-linux:
@@ -229,6 +242,14 @@ jobs:
           # since they'll just get checked in a separate shellcheck invocation.
           exclude: SC1091
 
+      - run:
+          name: ipynb-check
+          command: ./ci/clean_notebooks.py --check ./docs/tutorials
+
+      - run:
+          name: typeignore-check
+          command: ./ci/check_typeignore.py ${SRC_FILES}
+
       - run:
           name: flake8
           command: flake8 --version && flake8 -j "${NUM_CPUS}" ${SRC_FILES}
@@ -263,6 +284,10 @@ jobs:
           name: pytype
           command: pytype --version && pytype -j "${NUM_CPUS}" ${SRC_FILES[@]}
 
+      - run:
+          name: mypy
+          command: mypy --version && mypy ${SRC_FILES[@]} --exclude "${EXCLUDE_MYPY}" --follow-imports=silent --show-error-codes
+
   unit-test-linux:
     executor: unit-test-linux
     steps:

.coveragerc

@@ -0,0 +1,6 @@
+[report]
+exclude_lines =
+    pragma: no cover
+    @overload
+    @typing.overload
+    raise NotImplementedError

ci/Xdummy-entrypoint.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 
-# This script starts an X server and sets DISPLAY, then runs wrapped command.
+"""This script starts an X server and sets DISPLAY, then runs wrapped command."""
+
 # Usage: ./Xdummy-entrypoint.py [command]
 #
 # Adapted from https://github.com/openai/mujoco-py/blob/master/vendor/Xdummy-entrypoint
@@ -31,7 +32,7 @@
             "-config",
             "/etc/dummy_xorg.conf",
             ":0",
-        ]
+        ],
     )
     os.environ["DISPLAY"] = ":0"
 

ci/check_typeignore.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+
+"""Check for invalid "# type: ignore" comments.
+
+This script checks that no files in our source code have a "#type: ignore" comment
+without explicitly indicating the reason for the ignore. This is to ensure that we
+don't accidentally ignore errors that we should be fixing.
+"""
+import argparse
+import os
+import pathlib
+import re
+import sys
+from typing import List
+
+# Regex to match a "# type: ignore" comment not followed by a reason.
+TYPE_IGNORE_COMMENT = re.compile(r"#\s*type:\s*ignore\s*(?![^\[]*\[)")
+
+# Regex to match a "# type: ignore[<reason>]" comment.
+TYPE_IGNORE_REASON_COMMENT = re.compile(r"#\s*type:\s*ignore\[(?P<reason>.*)\]")
+
+
+class InvalidTypeIgnore(ValueError):
+    """Raised when a file has an invalid "# type: ignore" comment."""
+
+
+def check_file(file: pathlib.Path):
+    """Checks that the given file has no "# type: ignore" comments without a reason."""
+    with open(file, "r") as f:
+        for i, line in enumerate(f):
+            if TYPE_IGNORE_COMMENT.search(line):
+                raise InvalidTypeIgnore(
+                    f"{file}:{i+1}: Found a '# type: ignore' comment without a reason.",
+                )
+
+            if search := TYPE_IGNORE_REASON_COMMENT.search(line):
+                reason = search.group("reason")
+                if reason == "":
+                    raise InvalidTypeIgnore(
+                        f"{file}:{i+1}: Found a '# type: ignore[]' "
+                        "comment without a reason.",
+                    )
+
+
+def check_files(files: List[pathlib.Path]):
+    """Checks that the given files have no type: ignore comments without a reason."""
+    for file in files:
+        if file == pathlib.Path(__file__):
+            continue
+        check_file(file)
+
+
+def get_files_to_check(root_dirs: List[pathlib.Path]) -> List[pathlib.Path]:
+    """Returns a list of files that should be checked for "# type: ignore" comments."""
+    # Get the list of files that should be checked.
+    files = []
+    for root_dir in root_dirs:
+        for root, _, filenames in os.walk(root_dir):
+            for filename in filenames:
+                if filename.endswith(".py"):
+                    files.append(pathlib.Path(root) / filename)
+
+    return files
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "files",
+        nargs="+",
+        type=pathlib.Path,
+        help="List of files or paths to check for invalid '# type: ignore' comments.",
+    )
+    args = parser.parse_args()
+    return parser, args
+
+
+def main():
+    """Check for invalid "# type: ignore" comments."""
+    parser, args = parse_args()
+    file_list = get_files_to_check(args.files)
+    try:
+        check_files(file_list)
+    except InvalidTypeIgnore as e:
+        print(e)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ci/clean_notebooks.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python
+"""Clean all notebooks in the repository."""
+import argparse
+import pathlib
+import sys
+import traceback
+from typing import Any, Dict, List
+
+import nbformat
+
+
+class UncleanNotebookError(Exception):
+    """Raised when a notebook is unclean."""
+
+
+markdown_structure: Dict[str, Dict[str, Any]] = {
+    "cell_type": {"do": "keep"},
+    "metadata": {"do": "constant", "value": dict()},
+    "source": {"do": "keep"},
+    "id": {"do": "keep"},
+}
+
+code_structure: Dict[str, Dict[str, Any]] = {
+    "cell_type": {"do": "keep"},
+    "metadata": {"do": "constant", "value": dict()},
+    "source": {"do": "keep"},
+    "outputs": {"do": "constant", "value": list()},
+    "execution_count": {"do": "constant", "value": None},
+    "id": {"do": "keep"},
+}
+
+structure: Dict[str, Dict[str, Dict[str, Any]]] = {
+    "markdown": markdown_structure,
+    "code": code_structure,
+}
+
+
+def clean_notebook(file: pathlib.Path, check_only=False) -> None:
+    """Clean an ipynb notebook.
+
+    "Cleaning" means removing all output and metadata, as well as any other unnecessary
+    or vendor-dependent information or fields, so that it can be committed to the
+    repository, and so that artificial diffs are not introduced when the notebook is
+    executed.
+
+    Args:
+        file: Path to the notebook to clean.
+        check_only: If True, only check if the notebook is clean, and raise an
+            exception if it is not. If False, clean the notebook in-place.
+
+    Raises:
+        UncleanNotebookError: If `check_only` is True and the notebook is not clean.
+            Message contains brief description of the reason for the failure.
+        ValueError: unknown cell structure action.
+    """
+    # Read the notebook
+    with open(file) as f:
+        nb = nbformat.read(f, as_version=4)
+
+    was_dirty = False
+
+    if check_only:
+        print(f"Checking {file}")
+
+    for cell in nb.cells:
+
+        # Remove empty cells
+        if cell["cell_type"] == "code" and not cell["source"]:
+            if check_only:
+                raise UncleanNotebookError(f"Notebook {file} has empty code cell")
+            nb.cells.remove(cell)
+            was_dirty = True
+
+        # Clean the cell
+        # (copy the cell keys list so we can iterate over it while modifying it)
+        for key in list(cell):
+            if key not in structure[cell["cell_type"]]:
+                if check_only:
+                    raise UncleanNotebookError(
+                        f"Notebook {file} has unknown cell key {key}",
+                    )
+                del cell[key]
+                was_dirty = True
+            else:
+                cell_structure = structure[cell["cell_type"]][key]
+                if cell_structure["do"] == "keep":
+                    continue
+                elif cell_structure["do"] == "constant":
+                    constant_value = cell_structure["value"]
+                    if cell[key] != constant_value:
+                        if check_only:
+                            raise UncleanNotebookError(
+                                f"Notebook {file} has illegal cell value for key {key}"
+                                f" (value: {cell[key]}, "
+                                f"expected: {constant_value})",
+                            )
+                        cell[key] = constant_value
+                        was_dirty = True
+                else:
+                    raise ValueError(
+                        f"Unknown cell structure action {cell_structure['do']}",
+                    )
+
+    if not check_only and was_dirty:
+        # Write the notebook
+        with open(file, "w") as f:
+            nbformat.write(nb, f)
+        print(f"Cleaned {file}")
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    # if the argument --check has been passed, check if the notebooks are clean
+    # otherwise, clean them in-place
+    parser = argparse.ArgumentParser()
+    # capture files and paths to clean
+    parser.add_argument(
+        "files",
+        nargs="+",
+        type=pathlib.Path,
+        help="List of files or paths to clean",
+    )
+    parser.add_argument("--check", action="store_true")
+    args = parser.parse_args()
+    return parser, args
+
+
+def get_files(input_paths: List):
+    """Build list of files to scan from list of paths and files."""
+    files = []
+    for file in input_paths:
+        if file.is_dir():
+            files.extend(file.glob("**/*.ipynb"))
+        else:
+            if file.suffix == ".ipynb":
+                files.append(file)
+            else:
+                print(f"Skipping {file} (not a notebook)")
+    if not files:
+        print("No notebooks found")
+        sys.exit(1)
+    return files
+
+
+def main():
+    """Clean all notebooks in the repository, or check that they are clean."""
+    parser, args = parse_args()
+    check_only = args.check
+    input_paths = args.files
+
+    if len(input_paths) == 0:
+        parser.print_help()
+        sys.exit(1)
+
+    files = get_files(input_paths)
+
+    for file in files:
+        try:
+            clean_notebook(file, check_only=check_only)
+        except UncleanNotebookError:
+            traceback.print_exc()
+            sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ci/code_checks.sh

@@ -1,12 +1,25 @@
 #!/usr/bin/env bash
 
 # If you change these, also change .circleci/config.yml.
-SRC_FILES=(src/ tests/ experiments/ examples/ docs/conf.py setup.py)
+SRC_FILES=(src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/)
+EXCLUDE_MYPY="(?x)(
+  src/imitation/algorithms/preference_comparisons.py$
+  | src/imitation/rewards/reward_nets.py$
+  | src/imitation/util/sacred.py$
+  | src/imitation/algorithms/base.py$
+  | src/imitation/scripts/train_preference_comparisons.py$
+  | src/imitation/rewards/serialize.py$
+  | src/imitation/scripts/common/train.py$
+  | src/imitation/algorithms/mce_irl.py$
+  | src/imitation/algorithms/density.py$
+  | tests/algorithms/test_bc.py$
+)"
 
 set -x  # echo commands
 set -e  # quit immediately on error
 
 echo "Source format checking"
+./ci/clean_notebooks.py --check ./docs/tutorials/
 flake8 --darglint-ignore-regex '.*' "${SRC_FILES[@]}"
 black --check --diff "${SRC_FILES[@]}"
 codespell -I .codespell.skip --skip='*.pyc,tests/testdata/*,*.ipynb,*.csv' "${SRC_FILES[@]}"
@@ -22,6 +35,7 @@ fi
 if [ "${skipexpensive:-}" != "true" ]; then
   echo "Type checking"
   pytype -j auto "${SRC_FILES[@]}"
+  mypy "${SRC_FILES[@]}" --exclude "${EXCLUDE_MYPY}" --follow-imports=silent --show-error-codes
 
   echo "Building docs (validates docstrings)"
   pushd docs/