WIP: use pooch [ci skip]

Author: drammock

mne/datasets/utils.py

@@ -14,7 +14,9 @@
 import sys
 import zipfile
 import tempfile
+import pkg_resources
 from distutils.version import LooseVersion
+import pooch
 
 import numpy as np
 
@@ -222,87 +224,119 @@ def _data_path(path=None, force_update=False, update_path=True, download=True,
                name=None, check_version=False, return_version=False,
                archive_name=None, accept=False):
     """Aux function."""
-    key = {
-        'fake': 'MNE_DATASETS_FAKE_PATH',
-        'misc': 'MNE_DATASETS_MISC_PATH',
-        'sample': 'MNE_DATASETS_SAMPLE_PATH',
-        'spm': 'MNE_DATASETS_SPM_FACE_PATH',
-        'somato': 'MNE_DATASETS_SOMATO_PATH',
-        'brainstorm': 'MNE_DATASETS_BRAINSTORM_PATH',
-        'testing': 'MNE_DATASETS_TESTING_PATH',
-        'multimodal': 'MNE_DATASETS_MULTIMODAL_PATH',
-        'fnirs_motor': 'MNE_DATASETS_FNIRS_MOTOR_PATH',
-        'opm': 'MNE_DATASETS_OPM_PATH',
-        'visual_92_categories': 'MNE_DATASETS_VISUAL_92_CATEGORIES_PATH',
-        'kiloword': 'MNE_DATASETS_KILOWORD_PATH',
-        'mtrf': 'MNE_DATASETS_MTRF_PATH',
-        'fieldtrip_cmc': 'MNE_DATASETS_FIELDTRIP_CMC_PATH',
-        'phantom_4dbti': 'MNE_DATASETS_PHANTOM_4DBTI_PATH',
-        'limo': 'MNE_DATASETS_LIMO_PATH',
-        'refmeg_noise': 'MNE_DATASETS_REFMEG_NOISE_PATH',
-    }[name]
-
-    path = _get_path(path, key, name)
-    # To update the testing or misc dataset, push commits, then make a new
-    # release on GitHub. Then update the "releases" variable:
+    # To update the testing or misc datasets, push or merge commits to their
+    # respective repos, and make a new release of the dataset on GitHub.
+    # Then update the checksum in `mne/data/dataset_checksums.txt`,
+    # and change this "releases" variable:
     releases = dict(testing='0.112', misc='0.7')
-    # And also update the "md5_hashes['testing']" variable below.
-    # To update any other dataset, update the data archive itself (upload
-    # an updated version) and update the md5 hash.
-
-    # try to match url->archive_name->folder_name
-    urls = dict(  # the URLs to use
-        brainstorm=dict(
-            bst_auditory='https://osf.io/5t9n8/download?version=1',
-            bst_phantom_ctf='https://osf.io/sxr8y/download?version=1',
-            bst_phantom_elekta='https://osf.io/dpcku/download?version=1',
-            bst_raw='https://osf.io/9675n/download?version=2',
-            bst_resting='https://osf.io/m7bd3/download?version=3'),
-        fake='https://github.com/mne-tools/mne-testing-data/raw/master/'
-             'datasets/foo.tgz',
-        misc='https://codeload.github.com/mne-tools/mne-misc-data/'
-             'tar.gz/%s' % releases['misc'],
+    # To update any other dataset besides `testing` or `misc`, upload the new
+    # version of the data archive itself (e.g., to osf.io) and then update the
+    # corresponding checksum in `mne/data/dataset_checksums.txt`.
+    testing_data_name = f'mne-testing-data-{releases["testing"]}'
+    misc_data_name = f'mne-misc-data-{releases["misc"]}'
+
+    config_keys = dict(
+        fake='MNE_DATASETS_FAKE_PATH',
+        misc='MNE_DATASETS_MISC_PATH',
+        sample='MNE_DATASETS_SAMPLE_PATH',
+        spm='MNE_DATASETS_SPM_FACE_PATH',
+        somato='MNE_DATASETS_SOMATO_PATH',
+        brainstorm='MNE_DATASETS_BRAINSTORM_PATH',
+        testing='MNE_DATASETS_TESTING_PATH',
+        multimodal='MNE_DATASETS_MULTIMODAL_PATH',
+        fnirs_motor='MNE_DATASETS_FNIRS_MOTOR_PATH',
+        opm='MNE_DATASETS_OPM_PATH',
+        visual_92_categories='MNE_DATASETS_VISUAL_92_CATEGORIES_PATH',
+        kiloword='MNE_DATASETS_KILOWORD_PATH',
+        mtrf='MNE_DATASETS_MTRF_PATH',
+        fieldtrip_cmc='MNE_DATASETS_FIELDTRIP_CMC_PATH',
+        phantom_4dbti='MNE_DATASETS_PHANTOM_4DBTI_PATH',
+        limo='MNE_DATASETS_LIMO_PATH',
+        refmeg_noise='MNE_DATASETS_REFMEG_NOISE_PATH',
+    )
+    path = _get_path(path, config_keys[name], name)
+
+    # the download URLs
+    urls = dict(
+        bst_auditory='https://osf.io/5t9n8/download?version=1',
+        bst_phantom_ctf='https://osf.io/sxr8y/download?version=1',
+        bst_phantom_elekta='https://osf.io/dpcku/download?version=1',
+        bst_raw='https://osf.io/9675n/download?version=2',
+        bst_resting='https://osf.io/m7bd3/download?version=3',
+        fnirs_motor='https://osf.io/dj3eh/download?version=1',
+        kiloword='https://osf.io/qkvf9/download?version=1',
+        multimodal='https://ndownloader.figshare.com/files/5999598',
+        opm='https://osf.io/p6ae7/download?version=2',
+        phantom_4dbti='https://osf.io/v2brw/download?version=2',
         sample='https://osf.io/86qa2/download?version=5',
         somato='https://osf.io/tp4sg/download?version=7',
         spm='https://osf.io/je4s8/download?version=2',
-        testing='https://codeload.github.com/mne-tools/mne-testing-data/'
-                'tar.gz/%s' % releases['testing'],
-        multimodal='https://ndownloader.figshare.com/files/5999598',
-        fnirs_motor='https://osf.io/dj3eh/download?version=1',
-        opm='https://osf.io/p6ae7/download?version=2',
-        visual_92_categories=[
-            'https://osf.io/8ejrs/download?version=1',
-            'https://osf.io/t4yjp/download?version=1'],
+        visual_92_categories_1='https://osf.io/8ejrs/download?version=1',
+        visual_92_categories_2='https://osf.io/t4yjp/download?version=1',
         mtrf='https://osf.io/h85s2/download?version=1',
-        kiloword='https://osf.io/qkvf9/download?version=1',
-        fieldtrip_cmc='https://osf.io/j9b6s/download?version=1',
-        phantom_4dbti='https://osf.io/v2brw/download?version=2',
         refmeg_noise='https://osf.io/drt6v/download?version=1',
+        fieldtrip_cmc='https://osf.io/j9b6s/download?version=1',
+        fake=('https://github.com/mne-tools/mne-testing-data/raw/master/'
+              'datasets/foo.tgz'),
+        misc=('https://codeload.github.com/mne-tools/mne-misc-data/tar.gz/'
+              f'{releases["misc"]}'),
+        testing=('https://codeload.github.com/mne-tools/mne-testing-data/'
+                 f'tar.gz/{releases["testing"]}'),
     )
-    # filename of the resulting downloaded archive (only needed if the URL
-    # name does not match resulting filename)
+    # filename of the resulting downloaded archive
     archive_names = dict(
-        fieldtrip_cmc='SubjectCMC.zip',
+        bst_auditory='bst_auditory.tar.gz',
+        bst_phantom_ctf='bst_phantom_ctf.tar.gz',
+        bst_phantom_elekta='bst_phantom_elekta.tar.gz',
+        bst_raw='bst_raw.tar.gz',
+        bst_resting='bst_resting.tar.gz',
+        fnirs_motor='MNE-fNIRS-motor-data.tgz',
         kiloword='MNE-kiloword-data.tar.gz',
-        misc='mne-misc-data-%s.tar.gz' % releases['misc'],
-        mtrf='mTRF_1.5.zip',
         multimodal='MNE-multimodal-data.tar.gz',
-        fnirs_motor='MNE-fNIRS-motor-data.tgz',
         opm='MNE-OPM-data.tar.gz',
+        phantom_4dbti='MNE-phantom-4DBTi.zip',
         sample='MNE-sample-data-processed.tar.gz',
         somato='MNE-somato-data.tar.gz',
         spm='MNE-spm-face.tar.gz',
-        testing='mne-testing-data-%s.tar.gz' % releases['testing'],
-        visual_92_categories=['MNE-visual_92_categories-data-part1.tar.gz',
-                              'MNE-visual_92_categories-data-part2.tar.gz'],
-        phantom_4dbti='MNE-phantom-4DBTi.zip',
-        refmeg_noise='sample_reference_MEG_noise-raw.zip'
+        visual_92_categories_1='MNE-visual_92_categories-data-part1.tar.gz',
+        visual_92_categories_2='MNE-visual_92_categories-data-part2.tar.gz',
+        mtrf='mTRF_1.5.zip',
+        refmeg_noise='sample_reference_MEG_noise-raw.zip',
+        fieldtrip_cmc='SubjectCMC.zip',
+        fake='foo.tgz',
+        misc=f'{misc_data_name}.tar.gz',
+        testing=f'{testing_data_name}.tar.gz',
+    )
+    assert set(archive_names.keys()) == set(urls.keys())
+    # construct the mapping needed by pooch
+    pooch_urls = {archive_names[key]: urls[key] for key in urls}
+    # create the download manager
+    fetcher = pooch.create(
+        path=path,
+        base_url='',    # all URLs are given in the `urls` dict
+        version=None,   # because our data and code are in separate repos
+        registry=None,  # will load from file later
+        urls=pooch_urls
     )
+    # load the checksum registry
+    registry = pkg_resources.resource_stream(
+        'mne', op.join('data', 'dataset_checksums.txt'))
+    fetcher.load_registry(registry)
+    # update the keys that are versioned
+    versioned_keys = {
+        f'{testing_data_name}.tar.gz': fetcher.registry['mne-testing-data'],
+        f'{misc_data_name}.tar.gz': fetcher.registry['mne-misc-data']}
+    fetcher.registry.update(versioned_keys)
+    for key in ('testing', 'misc'):
+        del fetcher.registry[f'mne-{key}-data']
+
+    # TODO resume here
+
     # original folder names that get extracted (only needed if the
     # archive does not extract the right folder name; e.g., usually GitHub)
     folder_origs = dict(  # not listed means None (no need to move)
-        misc='mne-misc-data-%s' % releases['misc'],
-        testing='mne-testing-data-%s' % releases['testing'],
+        misc=misc_data_name,
+        testing=testing_data_name,
     )
     # finally, where we want them to extract to (only needed if the folder name
     # is not the same as the last bit of the archive name without the file
@@ -319,31 +353,7 @@ def _data_path(path=None, force_update=False, update_path=True, download=True,
         phantom_4dbti='MNE-phantom-4DBTi',
         refmeg_noise='MNE-refmeg-noise-data'
     )
-    md5_hashes = dict(
-        brainstorm=dict(
-            bst_auditory='fa371a889a5688258896bfa29dd1700b',
-            bst_phantom_ctf='80819cb7f5b92d1a5289db3fb6acb33c',
-            bst_phantom_elekta='1badccbe17998d18cc373526e86a7aaf',
-            bst_raw='fa2efaaec3f3d462b319bc24898f440c',
-            bst_resting='70fc7bf9c3b97c4f2eab6260ee4a0430'),
-        fake='3194e9f7b46039bb050a74f3e1ae9908',
-        misc='2b2f2fec9d1197ed459117db1c6341ee',
-        sample='12b75d1cb7df9dfb4ad73ed82f61094f',
-        somato='32fd2f6c8c7eb0784a1de6435273c48b',
-        spm='9f43f67150e3b694b523a21eb929ea75',
-        testing='8eabd73532dd7df7c155983962c5b1fd',
-        multimodal='26ec847ae9ab80f58f204d09e2c08367',
-        fnirs_motor='c4935d19ddab35422a69f3326a01fef8',
-        opm='370ad1dcfd5c47e029e692c85358a374',
-        visual_92_categories=['74f50bbeb65740903eadc229c9fa759f',
-                              '203410a98afc9df9ae8ba9f933370e20'],
-        kiloword='3a124170795abbd2e48aae8727e719a8',
-        mtrf='273a390ebbc48da2c3184b01a82e4636',
-        fieldtrip_cmc='6f9fd6520f9a66e20994423808d2528c',
-        phantom_4dbti='938a601440f3ffa780d20a17bae039ff',
-        refmeg_noise='779fecd890d98b73a4832e717d7c7c45'
-    )
-    assert set(md5_hashes.keys()) == set(urls.keys())
+
     url = urls[name]
     hash_ = md5_hashes[name]
     folder_orig = folder_origs.get(name, None)