Rename class names in type forward references (python-trio#53)

florimondmanca · sethmlarson · commit 2c2c9b7fdfd3 · 2020-01-04T19:10:00.000-06:00
* Unasync any string, test using forward ref to class itself

* Honor initial formatting
diff --git a/src/unasync/__init__.py b/src/unasync/__init__.py
@@ -60,6 +60,19 @@ def tokenize(f):
             last_end = (tok.end[0] + 1, 0)
 
 
+def unasync_name(name):
+    if name in ASYNC_TO_SYNC:
+        return ASYNC_TO_SYNC[name]
+    # Convert classes prefixed with 'Async' into 'Sync'
+    elif (
+        len(name) > 5
+        and name.startswith("Async")
+        and name[5].isupper()
+    ):
+        return "Sync" + name[5:]
+    return name
+
+
 def unasync_tokens(tokens):
     # TODO __await__, ...?
     used_space = None
@@ -72,15 +85,10 @@ def unasync_tokens(tokens):
             used_space = space
         else:
             if toknum == std_tokenize.NAME:
-                if tokval in ASYNC_TO_SYNC:
-                    tokval = ASYNC_TO_SYNC[tokval]
-                # Convert classes prefixed with 'Async' into 'Sync'
-                elif (
-                    len(tokval) > 5
-                    and tokval.startswith("Async")
-                    and tokval[5].isupper()
-                ):
-                    tokval = "Sync" + tokval[5:]
+                tokval = unasync_name(tokval)
+            elif toknum == std_tokenize.STRING:
+                left_quote, name, right_quote = tokval[0], tokval[1:-1], tokval[-1]
+                tokval = left_quote + unasync_name(name) + right_quote
             if used_space is None:
                 used_space = space
             yield (used_space, tokval)
diff --git a/tests/data/async/classes.py b/tests/data/async/classes.py
@@ -1,3 +1,13 @@
+class AsyncLock(object):
+    ...
+
+
 class AsyncSocket(object):
+    def __init__(self, send_lock: AsyncLock):
+        ...
+
+    async def __aenter__(self) -> "AsyncSocket":
+        ...
+
     async def send_all(self, data):
         ...
diff --git a/tests/data/sync/classes.py b/tests/data/sync/classes.py
@@ -1,3 +1,13 @@
+class SyncLock(object):
+    ...
+
+
 class SyncSocket(object):
+    def __init__(self, send_lock: SyncLock):
+        ...
+
+    def __enter__(self) -> "SyncSocket":
+        ...
+
     def send_all(self, data):
         ...
