RF: Try to import zstd from compression.zstd, backports.zstd, or pyzstd in

that order. Use module name "zstd" instead of "pyzstd"

Author: pauldmccarthy

nibabel/_compression.py

@@ -20,14 +20,20 @@
     import io
 
     import indexed_gzip  # type: ignore[import]
-    import pyzstd
+
+    # >= py314
+    try:
+        from compression import zstd  # type: ignore[import]
+    # < py314
+    except ImportError:
+        from backports import zstd  # type: ignore[import]
 
     HAVE_INDEXED_GZIP = True
     HAVE_ZSTD = True
 else:
     indexed_gzip, HAVE_INDEXED_GZIP, _ = optional_package('indexed_gzip')
-    pyzstd, HAVE_ZSTD, _ = optional_package('pyzstd')
-
+    zstd, HAVE_ZSTD, _ = optional_package(('compression.zstd',
+                                           'backports.zstd', 'pyzstd'))
 
 # Collections of types for isinstance or exception matching
 COMPRESSED_FILE_LIKES: tuple[type[io.IOBase], ...] = (
@@ -47,5 +53,5 @@
     IndexedGzipFile = gzip.GzipFile
 
 if HAVE_ZSTD:
-    COMPRESSED_FILE_LIKES += (pyzstd.ZstdFile,)
-    COMPRESSION_ERRORS += (pyzstd.ZstdError,)
+    COMPRESSED_FILE_LIKES += (zstd.ZstdFile,)
+    COMPRESSION_ERRORS += (zstd.ZstdError,)

nibabel/openers.py

@@ -16,7 +16,7 @@
 from bz2 import BZ2File
 from os.path import splitext
 
-from ._compression import HAVE_INDEXED_GZIP, IndexedGzipFile, pyzstd
+from ._compression import HAVE_INDEXED_GZIP, IndexedGzipFile, zstd
 
 if ty.TYPE_CHECKING:
     from types import TracebackType
@@ -103,9 +103,9 @@ def _zstd_open(
     mode: Mode = 'r',
     *,
     level_or_option: int | dict | None = None,
-    zstd_dict: pyzstd.ZstdDict | None = None,
-) -> pyzstd.ZstdFile:
-    return pyzstd.ZstdFile(filename, mode, level_or_option=level_or_option, zstd_dict=zstd_dict)
+    zstd_dict: zstd.ZstdDict | None = None,
+) -> zstd.ZstdFile:
+    return zstd.ZstdFile(filename, mode, level_or_option=level_or_option, zstd_dict=zstd_dict)
 
 
 class Opener: