Add unasync: remove feature

src/unasync/__init__.py

@@ -1,5 +1,6 @@
 """Top-level package for unasync."""
 
+import ast
 import collections
 import errno
 import os
@@ -67,18 +68,42 @@ def _unasync_file(self, filepath):
         with open(filepath, "rb") as f:
             encoding, _ = std_tokenize.detect_encoding(f.readline)
 
-        with open(filepath, 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))
+        with open(filepath, "rt", encoding=encoding) as f:
+            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)
 
-    def _unasync_tokens(self, tokens):
         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

tests/data/async/removals.py

@@ -0,0 +1,39 @@
+# isort: skip_file
+# fmt: off
+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
+
+# fmt: on

tests/data/sync/removals.py

@@ -0,0 +1,26 @@
+# isort: skip_file
+# fmt: off
+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')
+
+# fmt: on

tests/test_unasync.py

@@ -3,6 +3,7 @@
 import os
 import shutil
 import subprocess
+import sys
 
 import pytest
 
@@ -35,7 +36,6 @@ def test_rule_on_short_path():
 
 @pytest.mark.parametrize("source_file", TEST_FILES)
 def test_unasync(tmpdir, source_file):
-
     rule = unasync.Rule(fromdir=ASYNC_DIR, todir=str(tmpdir))
     rule._unasync_file(os.path.join(ASYNC_DIR, source_file))
 
@@ -64,16 +64,15 @@ def test_unasync_files(tmpdir):
 
 
 def test_build_py_modules(tmpdir):
-
     source_modules_dir = os.path.join(TEST_DIR, "example_mod")
     mod_dir = str(tmpdir) + "/" + "example_mod"
     shutil.copytree(source_modules_dir, mod_dir)
 
     env = copy.copy(os.environ)
     env["PYTHONPATH"] = os.path.realpath(os.path.join(TEST_DIR, ".."))
-    subprocess.check_call(["python", "setup.py", "build"], cwd=mod_dir, env=env)
+    subprocess.check_call([sys.executable, "setup.py", "build"], cwd=mod_dir, env=env)
     # Calling it twice to test the "if not copied" branch
-    subprocess.check_call(["python", "setup.py", "build"], cwd=mod_dir, env=env)
+    subprocess.check_call([sys.executable, "setup.py", "build"], cwd=mod_dir, env=env)
 
     unasynced = os.path.join(mod_dir, "build/lib/_sync/some_file.py")
     tree_build_dir = list_files(mod_dir)
@@ -84,14 +83,13 @@ def test_build_py_modules(tmpdir):
 
 
 def test_build_py_packages(tmpdir):
-
     source_pkg_dir = os.path.join(TEST_DIR, "example_pkg")
     pkg_dir = str(tmpdir) + "/" + "example_pkg"
     shutil.copytree(source_pkg_dir, pkg_dir)
 
     env = copy.copy(os.environ)
     env["PYTHONPATH"] = os.path.realpath(os.path.join(TEST_DIR, ".."))
-    subprocess.check_call(["python", "setup.py", "build"], cwd=pkg_dir, env=env)
+    subprocess.check_call([sys.executable, "setup.py", "build"], cwd=pkg_dir, env=env)
 
     unasynced = os.path.join(pkg_dir, "build/lib/example_pkg/_sync/__init__.py")
 
@@ -101,14 +99,13 @@ def test_build_py_packages(tmpdir):
 
 
 def test_project_structure_after_build_py_packages(tmpdir):
-
     source_pkg_dir = os.path.join(TEST_DIR, "example_pkg")
     pkg_dir = str(tmpdir) + "/" + "example_pkg"
     shutil.copytree(source_pkg_dir, pkg_dir)
 
     env = copy.copy(os.environ)
     env["PYTHONPATH"] = os.path.realpath(os.path.join(TEST_DIR, ".."))
-    subprocess.check_call(["python", "setup.py", "build"], cwd=pkg_dir, env=env)
+    subprocess.check_call([sys.executable, "setup.py", "build"], cwd=pkg_dir, env=env)
 
     _async_dir_tree = list_files(
         os.path.join(source_pkg_dir, "src/example_pkg/_async/.")
@@ -121,14 +118,13 @@ def test_project_structure_after_build_py_packages(tmpdir):
 
 
 def test_project_structure_after_customized_build_py_packages(tmpdir):
-
     source_pkg_dir = os.path.join(TEST_DIR, "example_custom_pkg")
     pkg_dir = str(tmpdir) + "/" + "example_custom_pkg"
     shutil.copytree(source_pkg_dir, pkg_dir)
 
     env = copy.copy(os.environ)
     env["PYTHONPATH"] = os.path.realpath(os.path.join(TEST_DIR, ".."))
-    subprocess.check_call(["python", "setup.py", "build"], cwd=pkg_dir, env=env)
+    subprocess.check_call([sys.executable, "setup.py", "build"], cwd=pkg_dir, env=env)
 
     _async_dir_tree = list_files(os.path.join(source_pkg_dir, "src/ahip/."))
     unasynced_dir_path = os.path.join(pkg_dir, "build/lib/hip/.")