Replace # pytype: disable suppressions with # pyrefly: ignore

envlogger/converters/codec.py

@@ -428,7 +428,7 @@ def from_bigint(int_bytes):
     )
 
   if array is None:
-    return None  # pytype: disable=bad-return-type
+    return None  # pyrefly: ignore[bad-return]
 
   return np.reshape(array, shape)
 

envlogger/converters/codec_test.py

@@ -1497,7 +1497,7 @@ def test_encode_dict_int_keys(self):
     v1 = t1.values.add().datum
     v1.shape.dim.add().size = -438
     v1.values.int64_values.append(456)
-    self.assertEqual(codec.encode({123: np.int64(456)}), expected)  # pyrefly: ignore[bad-argument-type, bad-assignment]
+    self.assertEqual(codec.encode({123: np.int64(456)}), expected)  # pyrefly: ignore[bad-assignment]
 
   def test_decode_dict_int_keys(self):
     """Dict with Python int keys."""
@@ -1515,7 +1515,8 @@ def test_decode_dict_int_keys(self):
   def test_identity_dict_int_keys(self):
     """Dict with Python int keys."""
     self.assertEqual(
-        codec.decode(codec.encode({123: np.int64(456)})), {123: np.int64(456)})  # pyrefly: ignore[bad-argument-type, bad-assignment]
+        codec.decode(codec.encode({123: np.int64(456)})), {123: np.int64(456)}  # pyrefly: ignore[bad-assignment]
+    )
 
   def test_encode_dict_int64_keys(self):
     """Dict with Python int64 keys."""
@@ -1528,7 +1529,7 @@ def test_encode_dict_int64_keys(self):
     v1 = t1.values.add().datum
     v1.shape.dim.add().size = -438
     v1.values.int32_values.append(12345)
-    self.assertEqual(codec.encode({np.int64(1729): np.int32(12345)}), expected)  # pyrefly: ignore[bad-argument-type, bad-assignment]
+    self.assertEqual(codec.encode({np.int64(1729): np.int32(12345)}), expected)  # pyrefly: ignore[bad-assignment]
 
   def test_decode_dict_int64_keys(self):
     """Dict with Python int64 keys."""
@@ -1546,8 +1547,9 @@ def test_decode_dict_int64_keys(self):
   def test_identity_dict_int64_keys(self):
     """Dict with Python int keys."""
     self.assertEqual(
-        codec.decode(codec.encode({np.int64(1729): np.int32(12345)})),  # pyrefly: ignore[bad-argument-type, bad-assignment]
-        {np.int64(1729): np.int32(12345)})
+        codec.decode(codec.encode({np.int64(1729): np.int32(12345)})),  # pyrefly: ignore[bad-assignment]
+        {np.int64(1729): np.int32(12345)},
+    )
 
   def test_identity_dict_mixed_keytypes(self):
     """Dict with Python mixed key types."""

envlogger/converters/spec_codec.py

@@ -50,8 +50,10 @@ def decode_environment_specs(
 ) -> dict[str, specs.Array | None]:
   """Decodes all the specs of an environment."""
   if encoded_specs:
-    return {spec_name: decode(encoded_specs[spec_name])  # pytype: disable=bad-return-type  # always-use-return-annotations
-            for spec_name in _ENVIRONMENT_SPEC_NAMES}
+    return {  # pyrefly: ignore[bad-return]
+        spec_name: decode(encoded_specs[spec_name])
+        for spec_name in _ENVIRONMENT_SPEC_NAMES
+    }
   return {spec_name: None for spec_name in _ENVIRONMENT_SPEC_NAMES}
 
 

envlogger/environment_logger_test.py

@@ -504,7 +504,7 @@ def _check_data(r: reader.Reader):
     futures = []
     with concurrent.futures.ThreadPoolExecutor(max_workers=n) as executor:
       for copy in copies:
-        futures.append(executor.submit(lambda c=copy: _check_data(c)))  # pyrefly: ignore[missing-argument]
+        futures.append(executor.submit(lambda c=copy: _check_data(c)))
     for f in futures:
       f.result(timeout=5)  # Wait for up to 5 seconds.
 

envlogger/testing/catch_env.py

@@ -58,7 +58,7 @@ def __init__(self, rows=10, columns=5, seed=1):
   def reset(self):
     """Returns the first `TimeStep` of a new episode."""
     self._reset_next_step = False
-    self._ball_x = self._rng.randint(self._columns)  # pyrefly: ignore[bad-assignment]
+    self._ball_x = self._rng.randint(self._columns)
     self._ball_y = 0
     self._paddle_x = self._columns // 2
     return dm_env.restart(self._observation())
@@ -73,7 +73,7 @@ def step(self, action):
     self._paddle_x = np.clip(self._paddle_x + dx, 0, self._columns - 1)
 
     # Drop the ball.
-    self._ball_y += 1  # pyrefly: ignore[unsupported-operation]
+    self._ball_y += 1
 
     # Check for termination.
     if self._ball_y == self._paddle_y: