diff --git a/src/unasync/__init__.py b/src/unasync/__init__.py
index e669b9a..872a1ff 100644
--- 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
@@ -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
diff --git a/tests/data/async/removals.py b/tests/data/async/removals.py
new file mode 100644
index 0000000..5640722
--- /dev/null
+++ b/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
diff --git a/tests/data/sync/removals.py b/tests/data/sync/removals.py
new file mode 100644
index 0000000..260a79f
--- /dev/null
+++ b/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
diff --git a/tests/test_unasync.py b/tests/test_unasync.py
index 35e0c6c..073e3fd 100644
--- a/tests/test_unasync.py
+++ b/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/.")