Merge pull request #32 from dimaqq/incoming

Dep update and cleanup after PR merge

Author: dimaqq

.gitignore

@@ -7,3 +7,4 @@ __pycache__/
 /awaitwhat/*.so
 graph.dot
 graph.svg
+.vscode/

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
--   repo: https://github.com/ambv/black
-    rev: stable
+-   repo: https://github.com/psf/black
+    rev: 21.6b0
     hooks:
     - id: black
-      python_version: python3.8
+      python_version: python3.9

awaitwhat/blocker.py

@@ -2,6 +2,7 @@
 import random
 from . import sleep
 from .stack import task_get_stack
+from . import wait_for
 
 
 def blockers(task):
@@ -12,8 +13,12 @@ def blockers(task):
 
     stack = task_get_stack(task, None)
 
-    if len(stack) > 2 and sleep.mine(stack[-2]):
-        return [sleep.decode(stack[-2])]
+    if len(stack) > 2:
+        if sleep.mine(stack[-2]):
+            return [sleep.decode(stack[-2])]
+        elif wait_for.mine(stack[-2]):
+            status, awaitable = wait_for.decode(stack[-2])
+            return [status, awaitable]
 
     # asyncio.gather()
     try:

awaitwhat/dot.py

@@ -4,8 +4,10 @@
 
 
 def label(vertex):
-    """ Graph vertex label in dot format """
+    """Graph vertex label in dot format"""
     label = f"{vertex.name} {vertex.state or ''}\n{vertex.traceback or ''}"
+    if not label.endswith("\n"):
+        label += "\n"
     label = json.dumps(label).replace("\\n", r"\l")
     return f"[label={label}]"
 

awaitwhat/graph.py

@@ -17,9 +17,12 @@ def new(tasks) -> Tuple[Set[Vertex], List[Edge]]:
     edges = list()
 
     for who in tasks:
-        src = Vertex.new(who, current)
+        children = blockers(who)
+        src = Vertex.new(
+            who, current, extra_text=[c for c in children if isinstance(c, str)]
+        )
         vertices.add(src)
-        for what in blockers(who):
+        for what in children:
             dst = Vertex.new(what, current)
             vertices.add(dst)
             edges.append(Edge(src, dst))
@@ -41,7 +44,8 @@ def __eq__(self, other):
         return isinstance(other, Vertex) and self.task == other.task
 
     @classmethod
-    def new(cls, task, current):
+    def new(cls, task, current, *, extra_text=()):
+        extra = "\n".join(extra_text)
         if isinstance(task, asyncio.Task):
             buf = io.StringIO()
             task_print_stack(task, None, buf)
@@ -53,11 +57,14 @@ def new(cls, task, current):
                 state = "done"
             else:
                 state = "current" if task is current else "pending"
-            return cls(name, state, concise_stack_trace(buf.getvalue()), task)
+            traceback = "\n".join(
+                filter(None, (concise_stack_trace(buf.getvalue()), extra))
+            )
+            return cls(name, state, traceback, task)
         elif isinstance(task, asyncio.Future):
-            return cls("Future", None, None, task)
+            return cls("Future", None, extra, task)
         else:
-            return cls(str(task), None, None, task)
+            return cls(str(task), None, extra, task)
 
 
 @dataclass(frozen=True)

awaitwhat/utils.py

@@ -13,4 +13,4 @@ def clean(line):
             line = re.sub('[^"]*/lib/python[0-9][.][0-9]/', "", line)
         return line
 
-    return '\n'.join(filter(None, (clean(line) for line in trace.split('\n'))))
+    return "\n".join(filter(None, (clean(line) for line in trace.split("\n"))))

awaitwhat/wait_for.py

@@ -0,0 +1,26 @@
+import asyncio
+
+
+def mine(frame):
+    return asyncio.wait_for.__code__ == frame.f_code
+
+
+def decode(frame):
+    try:
+        timeout = frame.f_locals["timeout"]
+    except Exception:
+        timeout = "?"
+    try:
+        timeout_handle_deadline = frame.f_locals["timeout_handle"].when()
+        now = frame.f_locals["waiter"].get_loop().time()
+        timeout_remaining = timeout_handle_deadline - now
+    except Exception:
+        timeout_remaining = "?"
+    try:
+        awaitable = frame.f_locals["fut"]
+    except Exception:
+        awaitable = None
+    return [
+        f"asyncio.wait_for: timeout {timeout} remaining {timeout_remaining}",
+        awaitable,
+    ]

build.py

@@ -1,4 +1,5 @@
 from distutils.core import Extension
 
+
 def build(setup_kwargs):
     setup_kwargs["ext_modules"] = [Extension("awaitwhat._what", ["awaitwhat/what.c"])]

examples/test_same_future.py

@@ -2,7 +2,8 @@
 import awaitwhat
 
 
-global_futures = list() # Structure to hold a global future
+global_futures = list()  # Structure to hold a global future
+
 
 async def main():
     # generate a global future

examples/test_wait_for.py

@@ -0,0 +1,48 @@
+import asyncio
+from examples.test_shield import frob_a_branch
+import awaitwhat
+
+
+FUTURES = []
+
+
+async def main():
+    t = asyncio.create_task(test())
+    await do_work()
+    await t
+
+
+async def do_work():
+    loop = asyncio.get_running_loop()
+    f = loop.create_future()
+    FUTURES.append(f)
+    branch = frob_a_branch(f)
+    await asyncio.wait_for(branch, 15)
+
+
+async def frob_a_branch(f):
+    await f
+
+
+def name(t):
+    try:
+        # Python3.8
+        return t.get_name()
+    except AttributeError:
+        return f"Task-{id(t)}"
+
+
+async def test():
+    import sys
+
+    await asyncio.sleep(1)
+    try:
+        tt = asyncio.all_tasks()
+        print(awaitwhat.dot.dumps(tt))
+    finally:
+        for f in FUTURES:
+            f.set_result(None)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())