Fix stderr suppression

pytorch-labs · Sep 20, 2024 · 123d1ee · 123d1ee
1 parent 4f2d18f
commit 123d1ee
Showing 1 changed file with 6 additions and 10 deletions.
diff --git a/torchfix/__main__.py b/torchfix/__main__.py
@@ -16,15 +16,6 @@
 from .common import CYAN, ENDC
 
 
-# Should get rid of this code eventually.
-@contextlib.contextmanager
-def StderrSilencer(redirect: bool = True):
-    if not redirect:
-        yield
-    with contextlib.redirect_stderr(io.StringIO()):
-        yield
-
-
 def _parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
 
@@ -87,7 +78,12 @@ def main() -> None:
     command_instance = TorchCodemod(codemod.CodemodContext(), config)
     DIFF_CONTEXT = 5
     try:
-        with StderrSilencer(not args.show_stderr):
+        supress_stderr = (
+            contextlib.redirect_stderr(io.StringIO())
+            if not args.show_stderr
+            else contextlib.nullcontext()
+        )
+        with supress_stderr:
             result = codemod.parallel_exec_transform_with_prettyprint(
                 command_instance,
                 torch_files,