Upgrade pytype (#801)

* Upgrade pytype and remove workaround for old versions

* new pytype need input directory or file

* fix np.dtype

* ignore typed-dict-error

* context manager related fix

* keep pytype checking more failures

* Move pytype config to pyproject.toml

* Use inputs specified in pyproject.toml

* Fix lint

* Fix undefined variable

* Fix end of file

* Fix typo

---------

Co-authored-by: Adam Gleave <[email protected]>

Author: ZiyueWang25

.pre-commit-config.yaml

@@ -78,7 +78,7 @@ repos:
     name: pytype
     language: system
     types: [python]
-    entry: "bash -c 'pytype -j ${NUM_CPUS:-auto}'"
+    entry: "bash -c 'pytype --keep-going -j ${NUM_CPUS:-auto}'"
     require_serial: true
     verbose: true
   - id: docs

pyproject.toml

@@ -14,3 +14,12 @@ build-backend = "setuptools.build_meta"
 
 [tool.black]
 target-version = ["py38"]
+
+[tool.pytype]
+inputs = [
+  "src/",
+  "tests/",
+  "experiments/",
+  "setup.py"
+]
+python_version = "3.8"

setup.cfg

@@ -59,11 +59,3 @@ omit =
 source =
     src/imitation
     *venv/lib/python*/site-packages/imitation
-
-[pytype]
-inputs =
-	src/
-	tests/
-	experiments/
-	setup.py
-python_version >= 3.8

setup.py

@@ -17,7 +17,7 @@
 ATARI_REQUIRE = [
     "seals[atari]~=0.2.1",
 ]
-PYTYPE = ["pytype==2022.7.26"] if IS_NOT_WINDOWS else []
+PYTYPE = ["pytype==2023.9.27"] if IS_NOT_WINDOWS else []
 
 # Note: the versions of the test and doc requirements should be tightly pinned to known
 #   working versions to make our CI/CD pipeline as stable as possible.

src/imitation/scripts/eval_policy.py

@@ -96,7 +96,7 @@ def eval_policy(
     sample_until = rollout.make_sample_until(eval_n_timesteps, eval_n_episodes)
     post_wrappers = [video_wrapper_factory(log_dir, **video_kwargs)] if videos else None
     render_mode = "rgb_array" if videos else None
-    with environment.make_venv(
+    with environment.make_venv(  # type: ignore[wrong-keyword-args]
         post_wrappers=post_wrappers,
         env_make_kwargs=dict(render_mode=render_mode),
     ) as venv:

src/imitation/scripts/ingredients/demonstrations.py

@@ -143,10 +143,10 @@ def _generate_expert_trajs(
         raise ValueError("n_expert_demos must be specified when generating demos.")
 
     logger.info(f"Generating {n_expert_demos} expert trajectories")
-    with environment.make_rollout_venv() as rollout_env:
+    with environment.make_rollout_venv() as env:  # type: ignore[wrong-arg-count]
         return rollout.rollout(
-            expert.get_expert_policy(rollout_env),
-            rollout_env,
+            expert.get_expert_policy(env),
+            env,
             rollout.make_sample_until(min_episodes=n_expert_demos),
             rng=_rnd,
         )

src/imitation/scripts/train_adversarial.py

@@ -119,7 +119,7 @@ def train_adversarial(
     custom_logger, log_dir = logging_ingredient.setup_logging()
     expert_trajs = demonstrations.get_expert_trajectories()
 
-    with environment.make_venv() as venv:
+    with environment.make_venv() as venv:  # type: ignore[wrong-arg-count]
         reward_net = reward.make_reward_net(venv)
         relabel_reward_fn = functools.partial(
             reward_net.predict_processed,

src/imitation/scripts/train_imitation.py

@@ -73,7 +73,7 @@ def bc(
     custom_logger, log_dir = logging_ingredient.setup_logging()
 
     expert_trajs = demonstrations.get_expert_trajectories()
-    with environment.make_venv() as venv:
+    with environment.make_venv() as venv:  # type: ignore[wrong-arg-count]
         bc_trainer = bc_ingredient.make_bc(venv, expert_trajs, custom_logger)
 
         bc_train_kwargs = dict(log_rollouts_venv=venv, **bc["train_kwargs"])
@@ -115,7 +115,7 @@ def dagger(
     if dagger["use_offline_rollouts"]:
         expert_trajs = demonstrations.get_expert_trajectories()
 
-    with environment.make_venv() as venv:
+    with environment.make_venv() as venv:  # type: ignore[wrong-arg-count]
         bc_trainer = bc_ingredient.make_bc(venv, expert_trajs, custom_logger)
 
         bc_train_kwargs = dict(log_rollouts_venv=venv, **bc["train_kwargs"])
@@ -161,7 +161,7 @@ def sqil(
     custom_logger, log_dir = logging_ingredient.setup_logging()
     expert_trajs = demonstrations.get_expert_trajectories()
 
-    with environment.make_venv() as venv:
+    with environment.make_venv() as venv:  # type: ignore[wrong-arg-count]
         sqil_trainer = sqil_algorithm.SQIL(
             venv=venv,
             demonstrations=expert_trajs,

src/imitation/scripts/train_preference_comparisons.py

@@ -166,7 +166,7 @@ def train_preference_comparisons(
 
     custom_logger, log_dir = logging_ingredient.setup_logging()
 
-    with environment.make_venv() as venv:
+    with environment.make_venv() as venv:  # type: ignore[wrong-arg-count]
         reward_net = reward.make_reward_net(venv)
         relabel_reward_fn = functools.partial(
             reward_net.predict_processed,

src/imitation/scripts/train_rl.py

@@ -99,7 +99,9 @@ def train_rl(
     policy_dir.mkdir(parents=True, exist_ok=True)
 
     post_wrappers = [lambda env, idx: wrappers.RolloutInfoWrapper(env)]
-    with environment.make_venv(post_wrappers=post_wrappers) as venv:
+    with environment.make_venv(  # type: ignore[wrong-keyword-args]
+        post_wrappers=post_wrappers,
+    ) as venv:
         callback_objs = []
         if reward_type is not None:
             reward_fn = load_reward(