diff --git a/src/unasync/__init__.py b/src/unasync/__init__.py index 148da41..634e08c 100644 --- 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 index 6969e1e..5a298ce 100644 --- 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 index 151a843..845e2ab 100644 --- 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): ...