Add post-processing hook for custom token transformations

src/unasync/__init__.py

@@ -69,7 +69,7 @@ def _unasync_file(self, filepath):
 
         with open(filepath, encoding=encoding) as f:
             tokens = tokenize_rt.src_to_tokens(f.read())
-            tokens = self._unasync_tokens(tokens)
+            tokens = self._transform_tokens(tokens)
             result = tokenize_rt.tokens_to_src(tokens)
             outfilepath = self.map_in_to_out_file_path(filepath)
             os.makedirs(os.path.dirname(outfilepath), exist_ok=True)
@@ -105,6 +105,25 @@ def _unasync_tokens(self, tokens):
 
                 yield token
 
+    def _transform_tokens(self, tokens):
+        """
+        Perform all token transformations.
+
+        The default implementation performs the standard async→sync
+        conversion. Subclasses may override this method to perform
+        additional token-level transformations.
+        """
+        tokens = self._unasync_tokens(tokens)
+        return self._postprocess_tokens(tokens)
+
+    def _postprocess_tokens(self, tokens):
+        """
+        Hook for subclasses.
+
+        Called after the standard async→sync conversion.
+        """
+        return tokens
+
     def unasync_name(self, name):
         if name in self.token_replacements:
             return self.token_replacements[name]

tests/data/postprocess/async/hello.py

@@ -0,0 +1,3 @@
+self._update_task = asyncio.current_task()
+asyncio
+current_task()

tests/data/postprocess/sync/hello.py

@@ -0,0 +1,3 @@
+self._update_task = current_thread()
+asyncio
+current_task()

tests/test_post_process.py

@@ -0,0 +1,63 @@
+import os
+
+from unasync import Rule
+
+TEST_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+TEST_DIR = os.path.join(TEST_DIR, "postprocess")
+ASYNC_DIR = os.path.join(TEST_DIR, "async")
+SYNC_DIR = os.path.join(TEST_DIR, "sync")
+TEST_FILES = sorted(f for f in os.listdir(ASYNC_DIR) if f.endswith(".py"))
+
+class PostProcessRule(Rule):
+
+    def _postprocess_tokens(self, tokens):
+        # Replace:
+        #     asyncio.current_task()
+        # with:
+        #     current_thread()
+
+        prev2 = None
+        prev1 = None
+
+        for token in tokens:
+
+            if (
+                prev2 is not None
+                and prev2.src == "asyncio"
+                and prev1.src == "."
+                and token.src == "current_task"
+            ):
+                yield token._replace(src="current_thread")
+
+                prev2 = None
+                prev1 = None
+
+            elif prev2 is not None:
+                yield prev2
+                prev2 = prev1
+                prev1 = token
+
+            else:
+                prev2 = prev1
+                prev1 = token
+
+        if prev2 is not None:
+            yield prev2
+        if prev1 is not None:
+            yield prev1
+
+
+def test_postprocess(tmpdir):
+    rule = PostProcessRule(fromdir=ASYNC_DIR, todir=str(tmpdir))
+
+    for source_file in TEST_FILES:
+        rule._unasync_file(os.path.join(ASYNC_DIR, source_file))
+
+    for source_file in TEST_FILES:
+        with open(os.path.join(SYNC_DIR, source_file)) as f:
+            truth = f.read()
+
+        with open(os.path.join(str(tmpdir), source_file)) as f:
+            unasynced = f.read()
+
+        assert unasynced == truth