possible bugfix indicated by mypy

florianziemen · Nov 15, 2024 · 3af99d3 · 3af99d3
1 parent 2bb8e79
commit 3af99d3
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/slkspec/xr_accessor.py b/slkspec/xr_accessor.py
@@ -69,7 +69,7 @@ def _get_output_tasks(
         return list(compress(k, layer_mask))
 
     def _get_output_keys(self, graph: dask.highlevelgraph.HighLevelGraph) -> set:
-        layer_keys = self._get_output_tasks(graph.layers)
+        layer_keys = self._get_output_tasks(graph)
         layers = [graph.layers[k] for k in layer_keys]
         output_keys = [lay.get_output_keys() for lay in layers]
 
@@ -119,9 +119,9 @@ def _get_data(self, data: da.Array) -> None:
         graph = data.dask.cull(keys=dask_keys)
 
         output_keys = self._get_output_keys(graph)
-        input_graph = self._get_input_graph(graph, output_keys)
-        input_graph = input_graph.to_dict()
-        input_graph = self._connect_tasks(input_graph, output_keys)
+        input_graph_dict = self._get_input_graph(graph, output_keys).to_dict()
+        input_graph = self._connect_tasks(input_graph_dict, output_keys)
+        del input_graph_dict
         scheduler = (
             dask.base.get_scheduler()
         )  # determine whether LocalCluster/SLUMCluster etc. exist