Add unasync: remove feature

Author: spyoungtech

src/unasync/__init__.py

@@ -2,11 +2,15 @@
 
 from __future__ import print_function
 
+import ast
 import collections
 import errno
 import os
 import sys
 import tokenize as std_tokenize
+from tokenize_rt import reversed_enumerate
+from tokenize_rt import src_to_tokens
+from tokenize_rt import tokens_to_src
 
 from setuptools.command import build_py as orig
 
@@ -74,11 +78,33 @@ def _unasync_file(self, filepath):
             tokens = _tokenize(f)
             tokens = self._unasync_tokens(tokens)
             result = _untokenize(tokens)
+            result = self._unasync_remove(contents=result, filename=filepath)
             outfilepath = filepath.replace(self.fromdir, self.todir)
             _makedirs_existok(os.path.dirname(outfilepath))
             with open(outfilepath, "w", **write_kwargs) as f:
                 print(result, file=f, end="")
 
+    def _unasync_remove(self, contents, filename):
+        tree = ast.parse(contents, filename=filename)
+        tokens = src_to_tokens(contents)
+        comment_lines_locations = []
+        for tok in tokens:
+            # find line numbers where "unasync: remove" comments are found
+            if tok.name == 'COMMENT' and 'unasync: remove' in tok.src:  # XXX: maybe make this a little more strict
+                comment_lines_locations.append(tok.line)
+        lines_to_remove = set()
+        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)
+        for i, tok in reversed_enumerate(tokens):
+            if tok.line in lines_to_remove:
+                tokens.pop(i)
+        new_contents = tokens_to_src(tokens)
+        return new_contents
+
     def _unasync_tokens(self, tokens):
         # TODO __await__, ...?
         used_space = None

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
+

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')
+