Fix type ignore errors

Author: Rocamonde

.circleci/config.yml

@@ -248,7 +248,7 @@ jobs:
 
       - run:
           name: typeignore-check
-          command: ./ci/check_typeignore.py
+          command: ./ci/check_typeignore.py ./src/ ./tests/
 
       - run:
           name: flake8

ci/check_typeignore.py

@@ -6,8 +6,7 @@
 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
@@ -51,29 +50,37 @@ def check_files(files: List[pathlib.Path]):
         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):
+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."""
-    files = get_files_to_check(pathlib.Path(root_dir))
+    parser, args = parse_args()
+    file_list = get_files_to_check(args.files)
     try:
-        check_files(files)
+        check_files(file_list)
     except InvalidTypeIgnore as e:
         print(e)
         sys.exit(1)

src/imitation/policies/replay_buffer_wrapper.py

@@ -23,7 +23,7 @@ def _samples_to_reward_fn_input(
     )
 
 
-class ReplayBufferRewardWrapper(BaseBuffer):
+class ReplayBufferRewardWrapper(ReplayBuffer):
     """Relabel the rewards in transitions sampled from a ReplayBuffer."""
 
     def __init__(

tests/algorithms/test_density_baselines.py

@@ -158,4 +158,4 @@ def test_density_trainer_raises(
         )
 
     with pytest.raises(TypeError, match="Unsupported demonstration type"):
-        density_trainer.set_demonstrations("foo")  # type: ignore
+        density_trainer.set_demonstrations("foo")

tests/data/test_rollout.py

@@ -300,7 +300,7 @@ def test_generate_trajectories_type_error(rng):
     sample_until = rollout.make_min_episodes(1)
     with pytest.raises(TypeError, match="Policy must be.*got <class 'str'> instead"):
         rollout.generate_trajectories(
-            "strings_are_not_valid_policies",  # type: ignore
+            "strings_are_not_valid_policies",
             venv,
             rng=rng,
             sample_until=sample_until,

tests/policies/test_replay_buffer_wrapper.py

@@ -38,17 +38,13 @@ def make_algo_with_wrapped_buffer(
         policy_kwargs=dict(),
         env=venv,
         seed=42,
-        # TODO(juan) we ignore the type below due to
-        #  https://github.com/DLR-RM/stable-baselines3/issues/1039
-        #  PR fixing this has been merged to master,
-        #  remove the type ignore in the next sb3 release.
-        replay_buffer_class=ReplayBufferRewardWrapper,  # type: ignore
+        replay_buffer_class=ReplayBufferRewardWrapper,
         replay_buffer_kwargs=dict(
             replay_buffer_class=replay_buffer_class,
             reward_fn=zero_reward_fn,
         ),
         buffer_size=buffer_size,
-    )  # type: ignore
+    )  # type: ignore[call-arg]
     return rl_algo
 
 
@@ -60,7 +56,7 @@ def test_invalid_args(rng):
         # we ignore the type because we are intentionally
         # passing the wrong type for the test
         make_algo_with_wrapped_buffer(
-            rl_cls=sb3.PPO,  # type: ignore
+            rl_cls=sb3.PPO,
             policy_cls=policies.ActorCriticPolicy,
             replay_buffer_class=buffers.ReplayBuffer,
             rng=rng,

tests/test_regularization.py

@@ -87,11 +87,11 @@ def test_interval_param_scaler_raises(interval_param_scaler):
     with pytest.raises(ValueError, match="train_loss must be a scalar"):
         scaler(1.0, th.Tensor([1.0, 2.0]), 1.0)
     with pytest.raises(ValueError, match="train_loss must be a scalar"):
-        scaler(1.0, "random value", th.tensor(1.0))  # type: ignore
+        scaler(1.0, "random value", th.tensor(1.0))
     with pytest.raises(ValueError, match="val_loss must be a scalar"):
-        scaler(1.0, 1.0, "random value")  # type: ignore
+        scaler(1.0, 1.0, "random value")
     with pytest.raises(ValueError, match="lambda_ must be a float"):
-        scaler(th.tensor(1.0), 1.0, 1.0)  # type: ignore
+        scaler(th.tensor(1.0), 1.0, 1.0)
     with pytest.raises(ValueError, match="lambda_ must not be zero.*"):
         scaler(0.0, 1.0, 1.0)
     with pytest.raises(ValueError, match="lambda_ must be non-negative.*"):
@@ -131,12 +131,12 @@ def test_interval_param_scaler_init_raises():
         ValueError,
         match="tolerable_interval must be a tuple of length 2",
     ):
-        updaters.IntervalParamScaler(0.5, (0.1, 0.9, 0.5))  # type: ignore
+        updaters.IntervalParamScaler(0.5, (0.1, 0.9, 0.5))
     with pytest.raises(
         ValueError,
         match="tolerable_interval must be a tuple of length 2",
     ):
-        updaters.IntervalParamScaler(0.5, (0.1,))  # type: ignore
+        updaters.IntervalParamScaler(0.5, (0.1,))
 
     # the first element of the interval must be at least 0.
     with pytest.raises(

tests/util/test_wb_logger.py

@@ -89,7 +89,7 @@ def finish(self):
 
 # we ignore the type below as one should technically not access the
 # __init__ method directly but only by creating an instance.
-@mock.patch.object(wandb, "__init__", mock_wandb.__init__)  # type: ignore
+@mock.patch.object(wandb, "__init__", mock_wandb.__init__)  # type: ignore[misc]
 @mock.patch.object(wandb, "init", mock_wandb.init)
 @mock.patch.object(wandb, "log", mock_wandb.log)
 @mock.patch.object(wandb, "finish", mock_wandb.finish)