Add 'unasync: remove' feature

python-trio · spyoungtech · May 4, 2023 · May 4, 2023 · May 4, 2023 · May 4, 2023
commit aaae7db7a30c7d9f09ab754c01bc3c619c0819c6
diff --git a/src/unasync/__init__.py b/src/unasync/__init__.py
@@ -1,5 +1,6 @@
 """Top-level package for unasync."""
 
+import ast
 import collections
 import errno
 import os
@@ -68,17 +69,37 @@ def _unasync_file(self, filepath):
             encoding, _ = std_tokenize.detect_encoding(f.readline)
 
         with open(filepath, "rt", encoding=encoding) as f:
-            tokens = tokenize_rt.src_to_tokens(f.read())
-            tokens = self._unasync_tokens(tokens)
-            result = tokenize_rt.tokens_to_src(tokens)
-            outfilepath = filepath.replace(self.fromdir, self.todir)
-            os.makedirs(os.path.dirname(outfilepath), exist_ok=True)
-            with open(outfilepath, "wb") as f:
-                f.write(result.encode(encoding))
-
-    def _unasync_tokens(self, tokens):
+            contents = f.read()
+        tokens = self._unasync_tokenize(contents=contents, filename=filepath)
+        result = tokenize_rt.tokens_to_src(tokens)
+        outfilepath = filepath.replace(self.fromdir, self.todir)
+        os.makedirs(os.path.dirname(outfilepath), exist_ok=True)
+        with open(outfilepath, "wb") as f:
+            f.write(result.encode(encoding))
+
+    def _unasync_tokenize(self, contents, filename):
+        tokens = tokenize_rt.src_to_tokens(contents)
+
+        comment_lines_locations = []
+        for token in tokens:
+            # find line numbers where "unasync: remove" comments are found
+            if token.name == 'COMMENT' and 'unasync: remove' in token.src:  # XXX: maybe make this a little more strict
+                comment_lines_locations.append(token.line)
+
+        lines_to_remove = set()
+        if comment_lines_locations:  # only parse ast if we actually have "unasync: remove" comments
+            tree = ast.parse(contents, filename=filename)
+            for node in ast.walk(tree):
+                # find nodes whose line number (start line) intersect with unasync: remove comments
+                if hasattr(node, 'lineno') and node.lineno in comment_lines_locations:
+                    for lineno in range(node.lineno, node.end_lineno + 1):
+                        # find all lines related to each node and mark those lines for removal
+                        lines_to_remove.add(lineno)
+
         skip_next = False
-        for i, token in enumerate(tokens):
+        for token in tokens:
+            if token.line in lines_to_remove:
+                continue
             if skip_next:
                 skip_next = False
                 continue

diff --git a/tests/data/async/removals.py b/tests/data/async/removals.py
@@ -0,0 +1,36 @@
+from common import (
+a, b , c  # these should stick around
+)
+
+# these imports should be removed
+from async_only import ( # unasync: remove
+    async_a, async_b,
+    async_c
+)
+
+CONST = 'foo'
+ASYNC_CONST = 'bar'  # unasync: remove
+
+async def foo():
+    print('this function should stick around')
+
+async def async_only(): # unasync: remove
+    print('this function will be removed entirely')
+
+
+class AsyncOnly: # unasync: remove
+    async def foo(self):
+        print('the entire class should be removed')
+
+
+class Foo:
+    async def foobar(self):
+        print('This method should stick around')
+
+    async def async_only_method(self): # unasync: remove
+        print('only this method should be removed')
+
+    async def another_method(self):
+        print('This line should stick around')
+        await self.something("the content in this line should be removed")  # unasync: remove
+
diff --git a/tests/data/sync/removals.py b/tests/data/sync/removals.py
@@ -0,0 +1,23 @@
+from common import (
+a, b , c  # these should stick around
+)
+
+# these imports should be removed
+
+CONST = 'foo'
+
+def foo():
+    print('this function should stick around')
+
+
+
+
+
+class Foo:
+    def foobar(self):
+        print('This method should stick around')
+
+
+    def another_method(self):
+        print('This line should stick around')
+