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 <[email protected]> * 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 <[email protected]> * Update tests/rewards/test_reward_fn.py Co-authored-by: Adam Gleave <[email protected]> * 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 <[email protected]> * Update src/imitation/util/util.py Co-authored-by: Adam Gleave <[email protected]> * 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 <[email protected]> * 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 <[email protected]> * Fix general type ignore in src * Fix type ignore errors * Formatting * Update src/imitation/util/util.py Co-authored-by: Adam Gleave <[email protected]> * Update src/imitation/scripts/common/rl.py Co-authored-by: Adam Gleave <[email protected]> * Update src/imitation/util/util.py Co-authored-by: Adam Gleave <[email protected]> * 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 <[email protected]>
HumanCompatibleAI · Oct 8, 2022 · 005c15f · 005c15f
1 parent a738955
commit 005c15f
Show file tree

Hide file tree

Showing 87 changed files with 2,039 additions and 693 deletions.
diff --git a/.circleci/config.yml b/.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:

diff --git a/.coveragerc b/.coveragerc
@@ -0,0 +1,6 @@
+[report]
+exclude_lines =
+    pragma: no cover
+    @overload
+    @typing.overload
+    raise NotImplementedError
diff --git a/ci/Xdummy-entrypoint.py b/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"
 

diff --git a/ci/check_typeignore.py b/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()
diff --git a/ci/clean_notebooks.py b/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()
diff --git a/ci/code_checks.sh b/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/