fix type hints

rapidsai · Feb 12, 2021 · f18137b · f18137b
1 parent 0751c4b
commit f18137b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/dask_cuda/proxify_device_objects.py b/dask_cuda/proxify_device_objects.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, List, MutableMapping
 
 from dask.utils import Dispatch
 
@@ -9,7 +9,7 @@
 
 def proxify_device_objects(
     obj: Any,
-    proxied_id_to_proxy: Dict[int, ProxyObject],
+    proxied_id_to_proxy: MutableMapping[int, ProxyObject],
     found_proxies: List[ProxyObject],
 ):
     """ Wrap device objects in ProxyObject
@@ -22,7 +22,7 @@ def proxify_device_objects(
     ----------
     obj: Any
         Object to search through or wrap in a ProxyObject.
-    proxied_id_to_proxy: Dict[int, ProxyObject]
+    proxied_id_to_proxy: MutableMapping[int, ProxyObject]
         Dict mapping the id() of proxied objects (CUDA device objects) to
         their proxy and is updated with all new proxied objects found in `obj`.
     found_proxies: List[ProxyObject]

diff --git a/dask_cuda/proxify_host_file.py b/dask_cuda/proxify_host_file.py
@@ -164,7 +164,7 @@ def get_dev_buffer_to_proxies(self) -> DefaultDict[Hashable, List[ProxyObject]]:
         with self.lock:
             # Notice, multiple proxy object can point to different non-overlapping
             # parts of the same device buffer.
-            ret = DefaultDict(list)
+            ret = defaultdict(list)
             for proxy in self.proxies_tally.get_unspilled_proxies():
                 for dev_buffer in proxy._obj_pxy_get_device_memory_objects():
                     ret[dev_buffer].append(proxy)