Resolve TODO after LibCST PR 994

pytorch-labs · Jan 25, 2024 · f751f1b · f751f1b
1 parent 9872f7f
commit f751f1b
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/torchfix/__main__.py b/torchfix/__main__.py
@@ -4,7 +4,6 @@
 import contextlib
 import sys
 import io
-import os
 
 from .torchfix import TorchCodemod, TorchCodemodConfig
 from .common import CYAN, ENDC
@@ -55,13 +54,11 @@ def main() -> None:
     MARKER = "torch"  # this will catch import torch or functorch
     torch_files = []
     for file in files:
-        # TODO: remove the check when https://github.com/Instagram/LibCST/pull/994 lands
-        if os.path.isfile(file):  # `codemod.gather_files` can return dirs with ".py"
-            with open(file, errors="replace") as f:
-                for line in f:
-                    if MARKER in line:
-                        torch_files.append(file)
-                        break
+        with open(file, errors="replace") as f:
+            for line in f:
+                if MARKER in line:
+                    torch_files.append(file)
+                    break
 
     config = TorchCodemodConfig()
     config.select = args.select