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

.circleci/config.yml

@@ -28,14 +28,14 @@ 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 ci/clean_notebooks.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 ci/clean_notebooks.py
+      SRC_FILES: src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/
       NUM_CPUS: 2
       EXCLUDE_MYPY: | 
         (?x)(
@@ -246,6 +246,10 @@ jobs:
           name: ipynb-check
           command: ./ci/clean_notebooks.py --check ./docs/tutorials
 
+      - run:
+          name: typeignore-check
+          command: ./ci/check_typeignore.py ./src/ ./tests/
+
       - run:
           name: flake8
           command: flake8 --version && flake8 -j "${NUM_CPUS}" ${SRC_FILES}

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

@@ -1,67 +1,89 @@
 #!/usr/bin/env python
 
-"""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."""
-
+"""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 re
 import pathlib
+import re
 import sys
 from typing import List
 
-import click
-
 # 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 ValueError(f"{file}:{i+1}: Found a '# type: ignore' comment without a reason.")
+                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 ValueError(f"{file}:{i+1}: Found a '# type: ignore[]' comment without a 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."""
+    """Checks that the given files have no type: ignore comments without a reason."""
     for file in files:
         check_file(file)
 
 
-def get_files_to_check(root_dir: pathlib.Path) -> List[pathlib.Path]:
+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, _, filenames in os.walk(root_dir):
-        for filename in filenames:
-            if filename.endswith(".py"):
-                files.append(pathlib.Path(root) / filename)
+    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
 
 
-@click.command()
-@click.option("--root-dir", type=click.Path(exists=True, file_okay=False, dir_okay=True), default="src")
-def main(root_dir: str):
-    """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."""
-    files = get_files_to_check(pathlib.Path(root_dir))
+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(files)
-    except ValueError as e:
+        check_files(file_list)
+    except InvalidTypeIgnore as e:
         print(e)
         sys.exit(1)
 
 
 if __name__ == "__main__":
-    main()
+    main()

ci/clean_notebooks.py

@@ -1,6 +1,10 @@
 #!/usr/bin/env python
 """Clean all notebooks in the repository."""
+import argparse
 import pathlib
+import sys
+import traceback
+from typing import Any, List, Tuple
 
 import nbformat
 
@@ -38,27 +42,28 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
     # Remove the output and metadata from each cell
     # also reset the execution count
     # if the cell has no code, remove it
+    fields_defaults: List[Tuple[str, Any]] = [
+        ("execution_count", None),
+        ("outputs", []),
+        ("metadata", {}),
+    ]
     for cell in nb.cells:
-        if "outputs" in cell and cell["outputs"]:
-            if check_only:
-                raise UncleanNotebookError(f"Notebook {file} has outputs")
-            cell["outputs"] = []
-            was_dirty = True
-        if "metadata" in cell and cell["metadata"]:
-            if check_only:
-                raise UncleanNotebookError(f"Notebook {file} has metadata")
-            cell["metadata"] = {}
-            was_dirty = True
-        if "execution_count" in cell and cell["execution_count"]:
-            if check_only:
-                raise UncleanNotebookError(f"Notebook {file} has execution count")
-            cell["execution_count"] = None
-            was_dirty = True
         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
+        for field, default in fields_defaults:
+            if cell.get(field) != default:
+                was_dirty = True
+                if check_only:
+                    raise UncleanNotebookError(
+                        f"Notebook {file} is not clean: cell has "
+                        f"field {field!r} with value {cell[field]!r} (expected "
+                        f"{default!r}). Cell:\n{cell['source']!r}",
+                    )
+                else:
+                    cell[field] = default
 
     if not check_only and was_dirty:
         # Write the notebook
@@ -67,12 +72,10 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
         print(f"Cleaned {file}")
 
 
-if __name__ == "__main__":
+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
-    import argparse
-    import traceback
-
     parser = argparse.ArgumentParser()
     # capture files and paths to clean
     parser.add_argument(
@@ -83,13 +86,13 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
     )
     parser.add_argument("--check", action="store_true")
     args = parser.parse_args()
-    check_only = args.check
-    # build list of files to scan from list of paths and files
+    return parser, args
+
+
+def get_files(input_paths: List):
+    """Build list of files to scan from list of paths and files."""
     files = []
-    if len(args.files) == 0:
-        parser.print_help()
-        exit(1)
-    for file in args.files:
+    for file in input_paths:
         if file.is_dir():
             files.extend(file.glob("**/*.ipynb"))
         else:
@@ -99,10 +102,29 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
                 print(f"Skipping {file} (not a notebook)")
     if not files:
         print("No notebooks found")
-        exit(1)
+        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()
-            exit(1)
+            sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ci/code_checks.sh

@@ -1,7 +1,19 @@
 #!/usr/bin/env bash
 
 # If you change these, also change .circleci/config.yml.
-SRC_FILES=(src/ tests/ experiments/ examples/ docs/conf.py setup.py ci/clean_notebooks.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
@@ -23,7 +35,7 @@ fi
 if [ "${skipexpensive:-}" != "true" ]; then
   echo "Type checking"
   pytype -j auto "${SRC_FILES[@]}"
-  mypy --show-error-codes "${SRC_FILES[@]}"
+  mypy "${SRC_FILES[@]}" --exclude "${EXCLUDE_MYPY}" --follow-imports=silent --show-error-codes
 
   echo "Building docs (validates docstrings)"
   pushd docs/

docs/tutorials/1_train_bc.ipynb

@@ -3,6 +3,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "[download this notebook here](https://github.com/HumanCompatibleAI/imitation/blob/master/docs/tutorials/1_train_bc.ipynb)\n",
     "# Train an Agent using Behavior Cloning\n",
@@ -16,6 +17,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "First we need some kind of expert in CartPole-v1 so we can sample some expert trajectories.\n",
     "For convenience we just train one using the stable-baselines3 library."
@@ -48,6 +50,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "Let's quickly check if the expert is any good.\n",
     "We usually should be able to reach a reward of 500, which is the maximum achievable value."
@@ -68,6 +71,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "Now we can use the expert to sample some trajectories.\n",
     "We flatten them right away since we are only interested in the individual transitions for behavior cloning.\n",
@@ -99,6 +103,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "Let's have a quick look at what we just generated using those library functions:"
    ]
@@ -120,6 +125,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "After we collected our transitions, it's time to set up our behavior cloning algorithm."
    ]
@@ -143,6 +149,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "As you can see the untrained policy only gets poor rewards:"
    ]
@@ -160,6 +167,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
+   "outputs": [],
    "source": [
     "After training, we can match the rewards of the expert (500):"
    ]
@@ -200,4 +208,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}