Suppressions for pyrefly

envlogger/backends/schedulers_test.py

@@ -123,12 +123,12 @@ def test_bernoulli_step_fixed_seed(self):
 
     # Run one trial with `seed`.
     scheduler = schedulers.BernoulliStepScheduler(
-        keep_probability=0.5, seed=seed)
+        keep_probability=0.5, seed=seed)  # pyrefly: ignore[bad-argument-type]
     outcomes = [scheduler(None) for _ in range(100)]
 
     # Repeat the trial with the same `seed`.
     other_scheduler = schedulers.BernoulliStepScheduler(
-        keep_probability=0.5, seed=seed)
+        keep_probability=0.5, seed=seed)  # pyrefly: ignore[bad-argument-type]
     other_outcomes = [other_scheduler(None) for _ in range(100)]
 
     # Assert that the outcomes are exactly the same.
@@ -240,15 +240,15 @@ def test_bernoulli_episode_fixed_seed(self):
 
     # Run one trial with `seed`.
     scheduler = schedulers.BernoulliEpisodeScheduler(
-        keep_probability=0.5, seed=seed)
+        keep_probability=0.5, seed=seed)  # pyrefly: ignore[bad-argument-type]
     outcomes = []
     for episode in episodes:
       for timestep in episode:
         outcomes.append(scheduler(timestep))
 
     # Repeat the trial with the same `seed`.
     other_scheduler = schedulers.BernoulliEpisodeScheduler(
-        keep_probability=0.5, seed=seed)
+        keep_probability=0.5, seed=seed)  # pyrefly: ignore[bad-argument-type]
     other_outcomes = []
     for episode in episodes:
       for timestep in episode:

envlogger/converters/codec_test.py

@@ -671,8 +671,8 @@ def test_decode_float32_list(self):
     datum.shape.dim.add().size = -438
     decoded = codec.decode(user_data)
     self.assertNotEmpty(decoded)
-    self.assertIsInstance(decoded[0], np.float32)
-    self.assertListEqual(decoded, [np.float32(3.14)])
+    self.assertIsInstance(decoded[0], np.float32)  # pyrefly: ignore[bad-index, unsupported-operation]
+    self.assertListEqual(decoded, [np.float32(3.14)])  # pyrefly: ignore[bad-argument-type]
 
   def test_encode_float32_nested_list(self):
     """Ensures that [[1.2, 3.4], [5.6, 7.8]] is represented correctly."""
@@ -717,8 +717,8 @@ def test_decode_float64_list(self):
     datum.shape.dim.add().size = -438
     decoded = codec.decode(user_data)
     self.assertNotEmpty(decoded)
-    self.assertIsInstance(decoded[0], np.float64)
-    self.assertListEqual(decoded, [np.float64(6.28)])
+    self.assertIsInstance(decoded[0], np.float64)  # pyrefly: ignore[bad-index, unsupported-operation]
+    self.assertListEqual(decoded, [np.float64(6.28)])  # pyrefly: ignore[bad-argument-type]
 
   ##############################################################################
   # int32
@@ -740,8 +740,8 @@ def test_decode_int32_list(self):
     datum.shape.dim.add().size = -438
     decoded = codec.decode(user_data)
     self.assertNotEmpty(decoded)
-    self.assertIsInstance(decoded[0], np.int32)
-    self.assertListEqual(decoded, [np.int32(-12345)])
+    self.assertIsInstance(decoded[0], np.int32)  # pyrefly: ignore[bad-index, unsupported-operation]
+    self.assertListEqual(decoded, [np.int32(-12345)])  # pyrefly: ignore[bad-argument-type]
 
   ##############################################################################
   # int64
@@ -763,8 +763,8 @@ def test_decode_int64_list(self):
     datum.shape.dim.add().size = -438
     decoded = codec.decode(user_data)
     self.assertNotEmpty(decoded)
-    self.assertIsInstance(decoded[0], np.int64)
-    self.assertListEqual(decoded, [np.int64(-1234567890123456)])
+    self.assertIsInstance(decoded[0], np.int64)  # pyrefly: ignore[bad-index, unsupported-operation]
+    self.assertListEqual(decoded, [np.int64(-1234567890123456)])  # pyrefly: ignore[bad-argument-type]
 
   # Homogeneity.
 
@@ -816,7 +816,7 @@ def test_decode_one_float_elem_ndarray_buffer(self):
     # 3.141519 in big-endian byte array.
     user_data.datum.values.float_values_buffer = b'\x40\x49\x0f\xd0'
     decoded = codec.decode(user_data)
-    self.assertEqual(decoded.dtype, np.float32)
+    self.assertEqual(decoded.dtype, np.float32)  # pyrefly: ignore[missing-attribute]
     np.testing.assert_equal(decoded, np.array([3.14159], dtype=np.float32))
 
   def test_encode_one_double_elem_scalar_ndarray(self):
@@ -1082,7 +1082,7 @@ def test_decode_one_int8_elem_ndarray(self):
     user_data.datum.shape.dim.add().size = 1
     user_data.datum.values.int8_values = b'\x91'
     decoded = codec.decode(user_data)
-    self.assertEqual(decoded.dtype, np.int8)
+    self.assertEqual(decoded.dtype, np.int8)  # pyrefly: ignore[missing-attribute]
     np.testing.assert_equal(decoded, np.array([-111], dtype=np.int8))
 
   def test_decode_two_int8_elem_ndarray(self):
@@ -1123,7 +1123,7 @@ def test_decode_one_int16_elem_ndarray(self):
     user_data.datum.shape.dim.add().size = 1
     user_data.datum.values.int16_values = b'\xfe\xa7'
     decoded = codec.decode(user_data)
-    self.assertEqual(decoded.dtype, np.int16)
+    self.assertEqual(decoded.dtype, np.int16)  # pyrefly: ignore[missing-attribute]
     np.testing.assert_equal(decoded, np.array([-345], dtype=np.int16))
 
   def test_decode_two_int16_elem_ndarray(self):
@@ -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)
+    self.assertEqual(codec.encode({123: np.int64(456)}), expected)  # pyrefly: ignore[bad-argument-type]
 
   def test_decode_dict_int_keys(self):
     """Dict with Python int keys."""
@@ -1515,7 +1515,7 @@ 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)})
+        codec.decode(codec.encode({123: np.int64(456)})), {123: np.int64(456)})  # pyrefly: ignore[bad-argument-type]
 
   def test_encode_dict_int64_keys(self):
     """Dict with Python int64 keys."""
@@ -1528,7 +1528,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)
+    self.assertEqual(codec.encode({np.int64(1729): np.int32(12345)}), expected)  # pyrefly: ignore[bad-argument-type]
 
   def test_decode_dict_int64_keys(self):
     """Dict with Python int64 keys."""
@@ -1546,13 +1546,13 @@ 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)})),
+        codec.decode(codec.encode({np.int64(1729): np.int32(12345)})),  # pyrefly: ignore[bad-argument-type]
         {np.int64(1729): np.int32(12345)})
 
   def test_identity_dict_mixed_keytypes(self):
     """Dict with Python mixed key types."""
     data = {123: np.int64(456), np.int64(1729): np.int32(12345), 'hello': True}
-    self.assertEqual(codec.decode(codec.encode(data)), data)
+    self.assertEqual(codec.decode(codec.encode(data)), data)  # pyrefly: ignore[bad-argument-type]
 
   ##############################################################################
   #

envlogger/converters/spec_codec_test.py

@@ -154,7 +154,7 @@ def test_encode_unsupported(self, arg):
   )
   def test_encode_array(self, input_spec, expected_spec_dict):
     """Checks that we can encode specs.Arrays."""
-    self._compare_spec_dicts(spec_codec.encode(input_spec), expected_spec_dict)
+    self._compare_spec_dicts(spec_codec.encode(input_spec), expected_spec_dict)  # pyrefly: ignore[bad-argument-type]
 
   # Single specs.BoundedArray.
 
@@ -237,7 +237,7 @@ def test_encode_array(self, input_spec, expected_spec_dict):
   def test_encode_bounded_array(self, input_spec, expected_spec_dict):
     """Checks that we can encode specs.BoundedArrays."""
     self._compare_spec_dicts(
-        spec_codec.encode(input_spec), expected_spec_dict)
+        spec_codec.encode(input_spec), expected_spec_dict)  # pyrefly: ignore[bad-argument-type]
 
   # Single specs.DiscreteArray.
 
@@ -260,7 +260,7 @@ def test_encode_bounded_array(self, input_spec, expected_spec_dict):
   def test_encode_discrete_array(self, input_spec, expected_spec_dict):
     """Checks that we can encode specs.DiscreArrays."""
     self._compare_spec_dicts(
-        spec_codec.encode(input_spec), expected_spec_dict)
+        spec_codec.encode(input_spec), expected_spec_dict)  # pyrefly: ignore[bad-argument-type]
 
   # Lists of specs.Arrays.
 
@@ -289,7 +289,7 @@ def test_encode_list_of_specs(self, input_spec, expected_spec_list):
     actual_spec_list = spec_codec.encode(input_spec)
     self.assertLen(actual_spec_list, len(expected_spec_list))
     for actual, expected in zip(actual_spec_list, expected_spec_list):
-      self._compare_spec_dicts(actual, expected)
+      self._compare_spec_dicts(actual, expected)  # pyrefly: ignore[bad-argument-type]
 
   # Tuples of specs.Arrays.
 
@@ -318,7 +318,7 @@ def test_encode_tuple_of_specs(self, input_spec, expected_spec_tuple):
     actual_spec_tuple = spec_codec.encode(input_spec)
     self.assertLen(actual_spec_tuple, len(expected_spec_tuple))
     for actual, expected in zip(actual_spec_tuple, expected_spec_tuple):
-      self._compare_spec_dicts(actual, expected)
+      self._compare_spec_dicts(actual, expected)  # pyrefly: ignore[bad-argument-type]
 
   # Dicts of specs.Arrays.
 

envlogger/environment_logger_test.py

@@ -400,7 +400,7 @@ def test_episode_starts_monotonically_increasing(self):
         previous = None
         for record in riegeli_reader.read_messages(storage_pb2.Datum):
           decoded = codec.decode_datum(record)
-          for episode_start, _ in decoded:
+          for episode_start, _ in decoded:  # pyrefly: ignore[not-iterable]
             if previous is None:
               continue
 
@@ -458,10 +458,10 @@ def only_first_step(data: step_data.StepData) -> bool:
 
     def my_episode_fn(timestep, unused_action, unused_env):
       if timestep.last():
-        my_episode_fn.x += 1
-      return my_episode_fn.x
+        my_episode_fn.x += 1  # pyrefly: ignore[missing-attribute]
+      return my_episode_fn.x  # pyrefly: ignore[missing-attribute]
 
-    my_episode_fn.x = 0
+    my_episode_fn.x = 0  # pyrefly: ignore[missing-attribute]
 
     with environment_logger.EnvLogger(
         catch_env.Catch(),
@@ -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)))
+        futures.append(executor.submit(lambda c=copy: _check_data(c)))  # pyrefly: ignore[missing-argument]
     for f in futures:
       f.result(timeout=5)  # Wait for up to 5 seconds.