Repository: dirty-cat/dirty_cat
Pull request #302 (2022-08-17 13:36:21)

Fetching: move `namedtuples` to `dataclasses`, and add target name to `DatasetInfoOnly`
Trying to use the no-loading version of the fetchers, I realized we're not passing the target name, and it's annoying to hard-code in downstream code, so I've added it to the `DatasetInfoOnly` requirements.
Also, since we're now in Python 3.8+ territory, I've figured we should use dataclasses instead of namedtuples, as their syntax is a little nicer (though more verbose), but most importantly, they integrate type hinting !

LilianBoulard: As this is a simple bugfix, I'm merging !

Patch:
diff --git a/dirty_cat/datasets/fetching.py b/dirty_cat/datasets/fetching.py
index 80291fd..1a818f6 100644
--- a/dirty_cat/datasets/fetching.py
+++ b/dirty_cat/datasets/fetching.py
@@ -22,20 +22,12 @@ import warnings
 import pandas as pd
 
 from pathlib import Path
-from collections import namedtuple
-from typing import Union, Dict, Any
+from dataclasses import dataclass
+from typing import Union, Dict, Any, List
 
 from dirty_cat.utils import Version
 from dirty_cat.datasets.utils import get_data_dir
 
-
-Details = namedtuple("Details", ["name", "file_id", "description"])
-Features = namedtuple("Features", ["names"])
-
-DatasetAll = namedtuple('Dataset', ('description', 'X', 'y', 'source', 'path'))
-DatasetInfoOnly = namedtuple('Dataset', ('description', 'source', 'path',
-                                         'read_csv_kwargs'))
-
 # Directory where the ``.gz`` files containing the
 # details on downloaded datasets are stored.
 # Note: the tree structure is created by ``fetch_openml()``.
@@ -63,6 +55,38 @@ TRAFFIC_VIOLATIONS_ID = 42132
 DRUG_DIRECTORY_ID = 43044
 
 
+@dataclass(unsafe_hash=True)
+class Details:
+    name: str
+    file_id: str
+    description: str
+
+
+@dataclass(unsafe_hash=True)
+class Features:
+    names: List[str]
+
+
+@dataclass(unsafe_hash=True)
+class DatasetAll:
+    name: str
+    description: str
+    X: pd.DataFrame
+    y: pd.Series
+    source: str
+    path: Path
+
+
+@dataclass(unsafe_hash=True)
+class DatasetInfoOnly:
+    name: str
+    description: str
+    source: str
+    target: str
+    path: Path
+    read_csv_kwargs: Dict[str, Any]
+
+
 def fetch_openml_dataset(dataset_id: int,
                          data_directory: Path = get_data_dir(),
                          ) -> Dict[str, Any]:
@@ -231,12 +255,11 @@ def _get_details(compressed_dir_path: Path) -> Details:
     # We filter out the irrelevant information.
     # If you want to modify this list (to add or remove items)
     # you must also modify the ``Details`` object definition.
-    f_details = {
-        "name": details["name"],
-        "file_id": details["file_id"],
-        "description": details["description"],
-    }
-    return Details(*f_details.values())
+    return Details(
+        name=details["name"],
+        file_id=details["file_id"],
+        description=details["description"],
+    )
 
 
 def _get_features(compressed_dir_path: Path) -> Features:
@@ -259,10 +282,9 @@ def _get_features(compressed_dir_path: Path) -> Features:
     # We filter out the irrelevant information.
     # If you want to modify this list (to add or remove items)
     # you must also modify the ``Features`` object definition.
-    features = {
-        "names": [column["name"] for column in raw_features["feature"]]
-    }
-    return Features(*features.values())
+    return Features(
+        names=[column["name"] for column in raw_features["feature"]]
+    )
 
 
 def _export_gz_data_to_csv(compressed_dir_path: Path,
@@ -300,10 +322,12 @@ def _features_to_csv_format(features: Features) -> str:
     return ",".join(features.names)
 
 
-def fetch_dataset_as_namedtuple(dataset_id: int, target: str,
-                                read_csv_kwargs: dict,
-                                load_dataframe: bool,
-                                ) -> Union[DatasetAll, DatasetInfoOnly]:
+def fetch_dataset_as_dataclass(dataset_name: str,
+                               dataset_id: int,
+                               target: str,
+                               read_csv_kwargs: dict,
+                               load_dataframe: bool,
+                               ) -> Union[DatasetAll, DatasetInfoOnly]:
     """
     Takes a dataset identifier, a target column name,
     and some additional keyword arguments for `pd.read_csv`.
@@ -326,6 +350,7 @@ def fetch_dataset_as_namedtuple(dataset_id: int, target: str,
         y = df[target]
         X = df.drop(target, axis='columns')
         dataset = DatasetAll(
+            name=dataset_name,
             description=info['description'],
             X=X,
             y=y,
@@ -334,8 +359,10 @@ def fetch_dataset_as_namedtuple(dataset_id: int, target: str,
         )
     else:
         dataset = DatasetInfoOnly(
+            name=dataset_name,
             description=info['description'],
             source=info['source'],
+            target=target,
             path=info['path'],
             read_csv_kwargs=read_csv_kwargs,
         )
@@ -376,12 +403,9 @@ def fetch_employee_salaries(load_dataframe: bool = True,
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    dataset = fetch_dataset_as_namedtuple(
+    dataset = fetch_dataset_as_dataclass(
+        dataset_name='Employee salaries',
         dataset_id=EMPLOYEE_SALARIES_ID,
         target='current_annual_salary',
         read_csv_kwargs={
@@ -419,12 +443,9 @@ def fetch_road_safety(load_dataframe: bool = True,
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    return fetch_dataset_as_namedtuple(
+    return fetch_dataset_as_dataclass(
+        dataset_name='Road safety',
         dataset_id=ROAD_SAFETY_ID,
         target='Sex_of_Driver',
         read_csv_kwargs={
@@ -456,12 +477,9 @@ def fetch_medical_charge(load_dataframe: bool = True
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    return fetch_dataset_as_namedtuple(
+    return fetch_dataset_as_dataclass(
+        dataset_name='Medical charge',
         dataset_id=MEDICAL_CHARGE_ID,
         target='Average_Total_Payments',
         read_csv_kwargs={
@@ -487,12 +505,9 @@ def fetch_midwest_survey(load_dataframe: bool = True
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    return fetch_dataset_as_namedtuple(
+    return fetch_dataset_as_dataclass(
+        dataset_name='Midwest survey',
         dataset_id=MIDWEST_SURVEY_ID,
         target='Census_Region',
         read_csv_kwargs={
@@ -519,12 +534,9 @@ def fetch_open_payments(load_dataframe: bool = True
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    return fetch_dataset_as_namedtuple(
+    return fetch_dataset_as_dataclass(
+        dataset_name='Open payments',
         dataset_id=OPEN_PAYMENTS_ID,
         target='status',
         read_csv_kwargs={
@@ -554,12 +566,9 @@ def fetch_traffic_violations(load_dataframe: bool = True
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    return fetch_dataset_as_namedtuple(
+    return fetch_dataset_as_dataclass(
+        dataset_name='Traffic violations',
         dataset_id=TRAFFIC_VIOLATIONS_ID,
         target='violation_type',
         read_csv_kwargs={
@@ -587,12 +596,9 @@ def fetch_drug_directory(load_dataframe: bool = True
 
     DatasetInfoOnly
         If `load_dataframe=False`
-
-    See Also
-    --------
-    dirty_cat.datasets.fetch_dataset_as_namedtuple : additional information
     """
-    return fetch_dataset_as_namedtuple(
+    return fetch_dataset_as_dataclass(
+        dataset_name='Drug directory',
         dataset_id=DRUG_DIRECTORY_ID,
         target='PRODUCTTYPENAME',
         read_csv_kwargs={


Test patch:
diff --git a/dirty_cat/datasets/tests/test_fetching.py b/dirty_cat/datasets/tests/test_fetching.py
index ac72826..24fa6af 100644
--- a/dirty_cat/datasets/tests/test_fetching.py
+++ b/dirty_cat/datasets/tests/test_fetching.py
@@ -131,10 +131,14 @@ def test_fetch_openml_dataset_mocked(mock_download, mock_export,
     we mock the functions to test its inner mechanisms.
     """
 
-    mock_get_details.return_value = Details("Dataset_name", "123456",
-                                            "Dummy dataset description.")
-    mock_get_features.return_value = Features(["id", "name", "transaction_id",
-                                               "owner", "recipient"])
+    mock_get_details.return_value = Details(
+        name="Dataset_name",
+        file_id="123456",
+        description="Dummy dataset description."
+    )
+    mock_get_features.return_value = Features(
+        names=["id", "name", "transaction_id", "owner", "recipient"]
+    )
 
     test_data_dir = get_test_data_dir()
 
@@ -296,19 +300,22 @@ def test__features_to_csv_format():
 
 
 @mock.patch('dirty_cat.datasets.fetching.fetch_openml_dataset')
-@mock.patch("dirty_cat.datasets.fetching.fetch_dataset_as_namedtuple")
-def test_import_all_datasets(mock_fetch_dataset_as_namedtuple,
+@mock.patch("dirty_cat.datasets.fetching.fetch_dataset_as_dataclass")
+def test_import_all_datasets(mock_fetch_dataset_as_dataclass,
                              mock_fetch_openml_dataset):
     """Tests functions ``fetch_*()``."""
 
     mock_fetch_openml_dataset.return_value = {
+        'name': 'Example dataset',
         'description': 'This is a dataset.',
         'source': 'https://www.openml.org/',
+        'target': 'To_predict',
         'path': Path("/path/to/file.csv"),
         'read_csv_kwargs': {'a': 'b'},
     }
 
     expected_return_value_all = fetching.DatasetAll(
+        name='Example dataset',
         description="This is a dataset.",
         source="https://www.openml.org/",
         path=Path("/path/to/file.csv"),
@@ -317,8 +324,10 @@ def test_import_all_datasets(mock_fetch_dataset_as_namedtuple,
     )
 
     expected_return_value_info_only = fetching.DatasetInfoOnly(
+        name='Example dataset',
         description="This is a dataset.",
         source="https://www.openml.org/",
+        target='To_predict',
         path=Path("/path/to/file.csv"),
         read_csv_kwargs={'a': 'b'},
     )
@@ -327,7 +336,7 @@ def test_import_all_datasets(mock_fetch_dataset_as_namedtuple,
         [expected_return_value_all, expected_return_value_info_only],
         [True, False],
     ):
-        mock_fetch_dataset_as_namedtuple.return_value = expected_return_value
+        mock_fetch_dataset_as_dataclass.return_value = expected_return_value
 
         returned_value = fetching.fetch_employee_salaries(drop_linked=False,
                                                           drop_irrelevant=False)