From 91e4df77a0b38f0c22541384e4fbc7987a24c406 Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jtomlinson@nvidia.com>
Date: Wed, 26 May 2021 15:36:47 +0100
Subject: [PATCH] Add HTML reprs for Client.who_has and Client.has_what

---
 distributed/client.py  |  5 ++--
 distributed/objects.py | 67 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 distributed/objects.py

diff --git a/distributed/client.py b/distributed/client.py
index 919d9181f2..245e53d0f9 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -53,6 +53,7 @@
 )
 from .diagnostics.plugin import UploadFile, WorkerPlugin, _get_worker_plugin_name
 from .metrics import time
+from .objects import HasWhat, WhoHas
 from .protocol import to_serialize
 from .protocol.pickle import dumps, loads
 from .publish import Datasets
@@ -3205,7 +3206,7 @@ def who_has(self, futures=None, **kwargs):
             keys = list(map(stringify, {f.key for f in futures}))
         else:
             keys = None
-        return self.sync(self.scheduler.who_has, keys=keys, **kwargs)
+        return WhoHas(self.sync(self.scheduler.who_has, keys=keys, **kwargs))
 
     def has_what(self, workers=None, **kwargs):
         """Which keys are held by which workers
@@ -3239,7 +3240,7 @@ def has_what(self, workers=None, **kwargs):
             workers = list(workers)
         if workers is not None and not isinstance(workers, (tuple, list, set)):
             workers = [workers]
-        return self.sync(self.scheduler.has_what, workers=workers, **kwargs)
+        return HasWhat(self.sync(self.scheduler.has_what, workers=workers, **kwargs))
 
     def processing(self, workers=None):
         """The tasks currently running on each worker
diff --git a/distributed/objects.py b/distributed/objects.py
new file mode 100644
index 0000000000..d53860514e
--- /dev/null
+++ b/distributed/objects.py
@@ -0,0 +1,67 @@
+"""This file contains custom objects.
+These are mostly regular objects with more useful _repr_ and _repr_html_ methods."""
+
+
+class HasWhat(dict):
+    """A dictionary of all workers and which keys that worker has."""
+
+    def _repr_html_(self):
+        rows = ""
+
+        for worker, keys in sorted(self.items()):
+            summary = ""
+            for key in keys:
+                summary += f"""<tr><td>{key}</td></tr>"""
+
+            rows += f"""<tr>
+            <td>{worker}</td>
+            <td>{len(keys)}</td>
+            <td>
+                <details>
+                <summary style='display:list-item'>Expand</summary>
+                <table>
+                {summary}
+                </table>
+                </details>
+            </td>
+        </tr>"""
+
+        output = f"""
+        <table>
+        <tr>
+            <th>Worker</th>
+            <th>Key count</th>
+            <th>Key list</th>
+        </tr>
+        {rows}
+        </table>
+        """
+
+        return output
+
+
+class WhoHas(dict):
+    """A dictionary of all keys and which workers have that key."""
+
+    def _repr_html_(self):
+        rows = ""
+
+        for title, keys in sorted(self.items()):
+            rows += f"""<tr>
+            <td>{title}</td>
+            <td>{len(keys)}</td>
+            <td>{", ".join(keys)}</td>
+        </tr>"""
+
+        output = f"""
+        <table>
+        <tr>
+            <th>Key</th>
+            <th>Copies</th>
+            <th>Workers</th>
+        </tr>
+        {rows}
+        </table>
+        """
+
+        return output