Move default location of data to the user home folder

CHANGES.rst

@@ -79,6 +79,9 @@ Minor changes
   by converting them to string before type inference.
   :pr:`623`by :user:`Leo Grinsztajn <LeoGrin>`
 
+* Moved the default storage location of data to the user's home folder.
+  :pr:`652` by :user:`Felix Lefebvre <flefebv>`
+
 Before skrub: dirty_cat
 ========================
 

skrub/datasets/_utils.py

@@ -2,7 +2,39 @@
 from pathlib import Path
 
 
-def get_data_dir(name: str | None = None) -> Path:
+def get_data_home(data_home=None) -> str:
+    """Returns the path of the skrub data directory.
+
+    This folder is used by some large dataset loaders to avoid downloading the
+    data several times.
+
+    By default the data directory is set to a folder named 'skrub_data' in the
+    user home folder.
+
+    Alternatively, it can be set programmatically by giving an explicit folder
+    path. The '~' symbol is expanded to the user home folder.
+
+    If the folder does not already exist, it is automatically created.
+
+    Parameters
+    ----------
+    data_home : str, default=None
+        The path to skrub data directory. If `None`, the default path
+        is `~/skrub_data`.
+
+    Returns
+    -------
+    data_home: str or path-like, default=None
+        The path to skrub data directory.
+    """
+    if data_home is None:
+        data_home = os.path.join("~", "skrub_data")
+    data_home = os.path.expanduser(data_home)
+    os.makedirs(data_home, exist_ok=True)
+    return data_home
+
+
+def get_data_dir(name: str | None = None, data_home: str | None = None) -> Path:
     """
     Returns the directory in which skrub looks for data.
 
@@ -11,13 +43,14 @@ def get_data_dir(name: str | None = None) -> Path:
 
     Parameters
     ----------
+    data_home : str, default=None
+        The path to skrub data directory. If `None`, the default path
+        is `~/skrub_data`.
     name: str, optional
         Subdirectory name. If omitted, the root data directory is returned.
     """
-    # Note: we stick to os.path instead of pathlib.Path because
-    # it's easier to test, and the functionality is the same.
-    module_path = Path(os.path.dirname(__file__)).resolve()
-    data_dir = module_path / "data"
+    data_home = get_data_home(data_home)
+    data_dir = Path(data_home).resolve()
     if name is not None:
         data_dir = data_dir / name
     return data_dir

skrub/datasets/tests/test_utils.py

@@ -1,18 +1,58 @@
+import os
+import shutil
+import tempfile
 from pathlib import Path
-from unittest import mock
 
+from skrub.datasets._utils import get_data_dir, get_data_home
 
-@mock.patch("os.path.dirname")
-def test_get_data_dir(mock_os_path_dirname):
+
+def test_get_data_dir():
     """
     Tests function ``get_data_dir()``.
     """
-    from skrub.datasets._utils import get_data_dir
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        expected_return_value_default = Path(tmpdirname)
+        assert get_data_dir(data_home=tmpdirname) == expected_return_value_default
+
+        expected_return_value_custom = expected_return_value_default / "tests"
+        assert (
+            get_data_dir(name="tests", data_home=tmpdirname)
+            == expected_return_value_custom
+        )
+
+
+def test_get_data_home_str():
+    """
+    Test function for ``get_data_home()`` when `data_home` is a string.
+    """
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # get_data_home will point to a pre-existing folder
+        assert os.path.exists(tmpdirname)
+        data_home = get_data_home(data_home=tmpdirname)
+        assert data_home == tmpdirname
+        assert os.path.exists(data_home)
+
+        # if the folder is missing it will be created again
+        shutil.rmtree(data_home)
+        assert not os.path.exists(data_home)
+        assert not os.path.exists(tmpdirname)
+        data_home = get_data_home(data_home=tmpdirname)
+        assert os.path.exists(data_home)
+
 
-    expected_return_value_default = Path("/user/directory/data").absolute()
+def test_get_data_home_None():
+    """
+    Test function for ``get_data_home()`` with `data_home` set to `None`.
+    """
+    # get path of the folder 'skrub_data' in the user home folder
+    user_path = os.path.join("~", "skrub_data")
+    user_path = os.path.expanduser(user_path)
 
-    mock_os_path_dirname.return_value = expected_return_value_default.parent
-    assert get_data_dir() == expected_return_value_default
+    # if the folder does not exist, create it, and clear it
+    if not os.path.exists(user_path):
+        data_home = get_data_home(data_home=None)
+        assert data_home == user_path
+        assert os.path.exists(data_home)
 
-    expected_return_value_custom = expected_return_value_default / "tests"
-    assert get_data_dir("tests") == expected_return_value_custom
+        shutil.rmtree(data_home)
+        assert not os.path.exists(data_home)