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/.")