Strenghten typechecking

The reason we were using `--protocols` was due to
google/pytype#764. Skipping just the offending
lines.

Author: mtrofin

.github/workflows/main.yml

@@ -24,7 +24,7 @@ jobs:
           pylint --version
       - name: Type check
         run: |
-          pytype --protocols -j auto .
+          pytype -j auto .
       - name: Test
         run: |
           pytest

compiler_opt/rl/compilation_runner_test.py

@@ -23,7 +23,9 @@
 from absl import flags
 import tensorflow as tf
 
-from google.protobuf import text_format
+# This is https://github.com/google/pytype/issues/764
+from google.protobuf import text_format  # pytype: disable=pyi-error
+
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import constant
 

compiler_opt/rl/data_collector_test.py

@@ -34,7 +34,7 @@ def test_build_distribution_monitor(self):
   def test_early_exit(self, mock_time):
     mock_time.return_value = 0
     early_exit = data_collector.EarlyExitChecker(
-        num_modules=10, deadline=10, thresholds=[(.9, 0), (.5, .5), (0, 1)])
+        num_modules=10, deadline=10, thresholds=((.9, 0), (.5, .5), (0, 1)))
 
     # We've waited no time, so have to hit 90% to early exit
     self.assertFalse(early_exit._should_exit(0))

compiler_opt/rl/local_data_collector_test.py

@@ -27,7 +27,9 @@
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import data_collector
 from compiler_opt.rl import local_data_collector
-from google.protobuf import text_format
+
+# This is https://github.com/google/pytype/issues/764
+from google.protobuf import text_format  # pytype: disable=pyi-error
 
 
 def _get_sequence_example(feature_value):

compiler_opt/rl/regalloc/regalloc_runner.py

@@ -24,7 +24,8 @@
 import gin
 import tensorflow as tf
 
-from google.protobuf import struct_pb2
+# This is https://github.com/google/pytype/issues/764
+from google.protobuf import struct_pb2  # pytype: disable=pyi-error
 from compiler_opt.rl import compilation_runner
 
 

compiler_opt/rl/train_locally.py

@@ -146,7 +146,7 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
                  len(reward_stat_map))
 
   data_collector = local_data_collector.LocalDataCollector(
-      file_paths=file_paths,
+      file_paths=tuple(file_paths),
       num_workers=FLAGS.num_workers,
       num_modules=FLAGS.num_modules,
       runner=runner,

compiler_opt/tools/__init__.py

@@ -96,7 +96,7 @@ def large_variation_counters(self, variation: float):
     return frozenset(ret)
 
   def _load_values(self, data: BenchmarkRunResults,
-                   names: Iterable[str]) -> PerBenchmarkResults:
+                   names: Iterable[str]) -> None:
     """Organize json values per-benchmark, per counter.
 
     Args:

compiler_opt/tools/benchmark_report.py

@@ -34,8 +34,7 @@
 import shutil
 import subprocess
 
-from typing import Dict
-from typing import List
+from typing import Dict, List, Optional
 
 from absl import app
 from absl import flags
@@ -71,15 +70,15 @@
 
 # TODO(ml-compiler-opt): maybe we can also convert here the cmdline file,from a
 # \0 - separated list of strings, to a \n one.
-def should_include_module(cmdline: str, match_regexp: str) -> bool:
+def should_include_module(cmdline: str, match_regexp: Optional[str]) -> bool:
   """Determine if the module should be included."""
   if match_regexp is None:
     return True
   lines = cmdline.split('\0')
   return any(len(re.findall(match_regexp, l)) for l in lines)
 
 
-def get_thinlto_index(cmdline: str, basedir: str) -> str:
+def get_thinlto_index(cmdline: str, basedir: str) -> Optional[str]:
   opts = cmdline.split('\0')
   for option in opts:
     if option.startswith('-fthinlto-index'):
@@ -150,7 +149,7 @@ def _get_extraction_bc_command(self, llvm_objcopy_path):
     ]
 
   def extract(self, llvm_objcopy_path: str, cmd_filter: str,
-              is_thinlto: bool) -> str:
+              is_thinlto: bool) -> Optional[str]:
     """Run llvm-objcopy to extract the .bc and command line."""
     if not os.path.exists(self.input_obj()):
       logging.info('%s does not exist.', self.input_obj())
@@ -235,7 +234,7 @@ def make_obj(obj_file: str) -> TrainingIRExtractor:
 
 # This is here just for readability, lint complains if the pooling expression is
 # over 3 lines; and it needs to be a non-local so it may be pickled.
-def extract_artifacts(obj: TrainingIRExtractor) -> str:
+def extract_artifacts(obj: TrainingIRExtractor) -> Optional[str]:
   return obj.extract(FLAGS.llvm_objcopy_path, FLAGS.cmd_filter,
                      FLAGS.thinlto_build)
 
@@ -272,7 +271,7 @@ def main(argv):
           f.write(path + '\n')
 
     logging.info('Converted %d files out of %d',
-                 len(objs) - relative_output_paths.count(None), len(objs))
+                 len(objs) - len(relative_output_paths), len(objs))
 
 
 if __name__ == '__main__':