From 2c9bf67a335c9a1e1f4d8316aa31cc43f659ca12 Mon Sep 17 00:00:00 2001
From: Simon Perkins <simon.perkins@gmail.com>
Date: Mon, 18 Mar 2024 17:27:57 +0200
Subject: [PATCH] Improve dask array name determinism (#10)

---
 HISTORY.rst              |  1 +
 xarrayfits/fits.py       | 31 ++-----------------------------
 xarrayfits/fits_proxy.py |  3 +++
 3 files changed, 6 insertions(+), 29 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index d9eec28..bea458e 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -4,6 +4,7 @@ History
 
 X.Y.Z (YYYY-MM-DD)
 ------------------
+* Improve dask array name determinism (:pr:`10`)
 * Change license from GPL3 to BSD3 (:pr:`9``)
 * Correct FITS Proxy Usage (:pr:`8`)
 * Update ruff settings (:pr:`7`)
diff --git a/xarrayfits/fits.py b/xarrayfits/fits.py
index 52def32..a1fb4e6 100644
--- a/xarrayfits/fits.py
+++ b/xarrayfits/fits.py
@@ -2,7 +2,7 @@
 
 """Main module."""
 
-from functools import partial, reduce
+from functools import reduce
 from itertools import product
 import logging
 import os
@@ -38,32 +38,6 @@ def short_fits_file(table_name):
     return os.path.split(table_name.rstrip(os.sep))[1]
 
 
-def fits_open_graph(fits_file, **kwargs):
-    """
-    Generate a dask graph containing fits open commands
-
-    Parameters
-    ----------
-    fits_file : str
-        FITS filename
-    **kwargs (optional) :
-        Keywords arguments passed to the :meth:`astropy.io.fits.open`
-        command.`
-
-    Returns
-    -------
-    tuple
-        Graph key associated with the opened file
-    dict
-        Dask graph containing the graph open command
-
-    """
-    token = dask.base.tokenize(fits_file, kwargs)
-    fits_key = ("open", short_fits_file(fits_file), token)
-    fits_graph = {fits_key: (partial(FitsProxy, **kwargs), fits_file)}
-    return fits_key, fits_graph
-
-
 BITPIX_MAP = {
     np.int8: 8,
     np.int16: 16,
@@ -109,9 +83,8 @@ def generate_slice_gets(fits_proxy, hdu, shape, dtype, chunks):
         with the ``hdu``.
     """
 
-    token = dask.base.tokenize(fits_proxy)
+    token = dask.base.tokenize(fits_proxy, hdu, dtype)
     name = "-".join((short_fits_file(fits_proxy._filename), "slice", token))
-
     dsk_chunks = da.core.normalize_chunks(chunks, shape)
 
     # Produce keys and slices
diff --git a/xarrayfits/fits_proxy.py b/xarrayfits/fits_proxy.py
index f9e67d0..d268330 100644
--- a/xarrayfits/fits_proxy.py
+++ b/xarrayfits/fits_proxy.py
@@ -57,5 +57,8 @@ def hdu(self):
                     self._hdul = fits.open(self._filename, **self._kwargs)
                     return self._hdul
 
+    def __hash__(self):
+        return hash((self._filename, tuple(set(self._kwargs.items()))))
+
     def __reduce__(self):
         return (FitsProxy.from_reduce_args, (self._filename, self._kwargs))