refactor(log): address too-many-statements, too-many-branches, and too-many-instance-attributes

imbens/datasets/_openml.py

@@ -21,6 +21,95 @@
 import tqdm
 
 
+def _get_data_path(data_home):
+    if data_home is None:
+        data_home = Path(user_cache_dir("imbens")) / "datasets"
+        data_home.mkdir(parents=True, exist_ok=True)
+        return data_home
+    if not isinstance(data_home, Path):
+        data_home = Path(data_home)
+    if not data_home.is_dir():
+        raise ValueError(f"data_home {data_home} is not a directory")
+    data_home.mkdir(parents=True, exist_ok=True)
+    return data_home
+
+
+def _save_data(X, y, data_home, file_name):
+    np.savez_compressed(data_home / file_name, X=X, y=y)
+
+
+def _load_data(data_home, file_name):
+    data = np.load(data_home / file_name, allow_pickle=True)
+    return data["X"], data["y"]
+
+
+def _validate_openml_params(openml_id, cat_preprocess, data_home):
+    assert isinstance(openml_id, int), "openml_id must be an integer"
+    assert cat_preprocess in [
+        "drop",
+        "onehot",
+        "ordinal",
+    ], "cat_preprocess must be one of ['drop', 'onehot', 'ordinal']"
+    assert isinstance(
+        data_home, (str, Path, type(None))
+    ), "data_home must be a string, Path or None"
+
+
+def _encode_target(y):
+    y_map = y.value_counts().sort_values()[::-1].index
+    y_map = {v: i for i, v in enumerate(y_map)}
+    return y.map(y_map)
+
+
+def _standardize_numeric(X, feats_type):
+    scaler = StandardScaler()
+    X[feats_type["num"]] = scaler.fit_transform(X[feats_type["num"]])
+    return X
+
+
+def _preprocess_feature(feat, X, feats_type):
+    if is_any_real_numeric_dtype(X[feat]):
+        feats_type["num"].append(feat)
+        if X[feat].isnull().sum() > 0:
+            X[feat] = X[feat].fillna(X[feat].mean())
+    else:
+        if X[feat].isnull().sum() > 0:
+            X[feat] = X[feat].fillna(X[feat].mode().iloc[0])
+        n_unique = len(X[feat].unique())
+        if n_unique > 2:
+            try:
+                X[feat] = pd.to_numeric(X[feat])
+            except ValueError:
+                if n_unique <= 50:
+                    feats_type["multi_cat"].append(feat)
+                else:
+                    feats_type["drop"].append(feat)
+        else:
+            feats_type["bin_cat"].append(feat)
+
+
+def _cat_drop(X, feats_type, ord_encoder):
+    return X.drop(columns=feats_type["multi_cat"] + feats_type["bin_cat"])
+
+
+def _cat_onehot(X, feats_type, ord_encoder):
+    X[feats_type["bin_cat"]] = ord_encoder.fit_transform(X[feats_type["bin_cat"]])
+    return pd.get_dummies(X, columns=feats_type["multi_cat"])
+
+
+def _cat_ordinal(X, feats_type, ord_encoder):
+    X[feats_type["bin_cat"]] = ord_encoder.fit_transform(X[feats_type["bin_cat"]])
+    X[feats_type["multi_cat"]] = ord_encoder.fit_transform(X[feats_type["multi_cat"]])
+    return X
+
+
+_CAT_PREPROCESSORS = {
+    "drop": _cat_drop,
+    "onehot": _cat_onehot,
+    "ordinal": _cat_ordinal,
+}
+
+
 def _fetch_openml_data(openml_id, cat_preprocess="onehot", data_home=None):
     """
     Fetches a dataset from OpenML, preprocesses it, and caches it locally.
@@ -52,112 +141,35 @@ def _fetch_openml_data(openml_id, cat_preprocess="onehot", data_home=None):
     If the dataset is already cached locally, it will be loaded from the cache.
     Otherwise, the dataset will be downloaded, preprocessed, and saved locally.
     """
+    _validate_openml_params(openml_id, cat_preprocess, data_home)
 
-    def get_data_path(data_home):
-        if data_home is None:
-            data_home = Path(user_cache_dir("imbens")) / "datasets"
-            data_home.mkdir(parents=True, exist_ok=True)
-            return data_home
-        else:
-            # check if store_path is a valid path
-            if not isinstance(data_home, Path):
-                data_home = Path(data_home)
-            if not data_home.is_dir():
-                raise ValueError(f"data_home {data_home} is not a directory")
-            data_home.mkdir(parents=True, exist_ok=True)
-            return data_home
-
-    def save_data(X, y, data_home, file_name):
-        np.savez_compressed(data_home / file_name, X=X, y=y)
-        # print(f"Data saved to {data_home / file_name}")
-
-    def load_data(data_home, file_name):
-        data = np.load(data_home / file_name, allow_pickle=True)
-        X = data["X"]
-        y = data["y"]
-        return X, y
-
-    assert isinstance(openml_id, int), "openml_id must be an integer"
-    assert cat_preprocess in [
-        "drop",
-        "onehot",
-        "ordinal",
-    ], "cat_preprocess must be one of ['drop', 'onehot', 'ordinal']"
-    assert isinstance(
-        data_home, (str, Path, type(None))
-    ), "data_home must be a string, Path or None"
-
-    data_home = get_data_path(data_home)
+    data_home = _get_data_path(data_home)
     dataset = openml.datasets.get_dataset(openml_id)
     file_name = f"{dataset.id}_{cat_preprocess}_{dataset.name}.npz"
 
-    # check if data is already cached
     try:
-        X, y = load_data(data_home, file_name)
-        # print(f"Data loaded from {data_home / file_name}")
+        X, y = _load_data(data_home, file_name)
         return X, y
     except FileNotFoundError:
-        # print(f"Data not cached in {data_home}, processing from scratch.")
         pass
 
     X, y, cat_ind, feat_names = dataset.get_data(
         target=dataset.default_target_attribute
     )
 
-    # target encoding, smaller classes aer assigned larger values
-    y_map = y.value_counts().sort_values()[::-1].index
-    y_map = {v: i for i, v in enumerate(y_map)}
-    y = y.map(y_map)
+    y = _encode_target(y)
 
     feats_type = {"num": [], "bin_cat": [], "multi_cat": [], "drop": []}
-    # preprocessing
     for feat in X.columns:
-        if is_any_real_numeric_dtype(X[feat]):
-            feats_type["num"].append(feat)
-            if X[feat].isnull().sum() > 0:
-                # for numerical columns, fill nan with mean
-                X[feat] = X[feat].fillna(X[feat].mean())
-        else:  # categorical column
-            if X[feat].isnull().sum() > 0:
-                # for categorical columns, fill nan with most frequent value
-                X[feat] = X[feat].fillna(X[feat].mode().iloc[0])
-            n_unique = len(X[feat].unique())
-            if n_unique > 2:
-                # try to convert to numeric
-                try:
-                    X[feat] = pd.to_numeric(X[feat])
-                except ValueError:
-                    if n_unique <= 50:
-                        feats_type["multi_cat"].append(feat)
-                    else:
-                        feats_type["drop"].append(feat)
-            else:
-                feats_type["bin_cat"].append(feat)
+        _preprocess_feature(feat, X, feats_type)
 
     ord_encoder = OrdinalEncoder()
     X = X.drop(columns=feats_type["drop"])
-    # encode categorical columns
-    if cat_preprocess == "drop":
-        X = X.drop(columns=feats_type["multi_cat"])
-        X = X.drop(columns=feats_type["bin_cat"])
-    elif cat_preprocess == "onehot":
-        X[feats_type["bin_cat"]] = ord_encoder.fit_transform(X[feats_type["bin_cat"]])
-        X = pd.get_dummies(X, columns=feats_type["multi_cat"])
-    elif cat_preprocess == "ordinal":
-        # ordinal encoding for multi categorical columns
-        X[feats_type["bin_cat"]] = ord_encoder.fit_transform(X[feats_type["bin_cat"]])
-        X[feats_type["multi_cat"]] = ord_encoder.fit_transform(
-            X[feats_type["multi_cat"]]
-        )
-    else:
-        raise ValueError(f"Unknown cat_preprocess: {cat_preprocess}")
+    X = _CAT_PREPROCESSORS[cat_preprocess](X, feats_type, ord_encoder)
 
-    # standardize numerical columns
-    scaler = StandardScaler()
-    X[feats_type["num"]] = scaler.fit_transform(X[feats_type["num"]])
+    X = _standardize_numeric(X, feats_type)
 
-    # save data
-    save_data(X, y, data_home, file_name)
+    _save_data(X, y, data_home, file_name)
 
     return X, y
 

imbens/datasets/_zenodo.py

@@ -100,6 +100,57 @@
     MAP_ID_NAME[v + 1] = k
 
 
+def _resolve_filter_data(filter_data):
+    """Validate and resolve filter_data into a list of dataset names."""
+    if filter_data is None:
+        return list(MAP_NAME_ID.keys())
+
+    list_data = MAP_NAME_ID.keys()
+    filter_data_ = []
+    for it in filter_data:
+        if isinstance(it, str):
+            if it not in list_data:
+                raise ValueError(
+                    f"{it} is not a dataset available. "
+                    f"The available datasets are {list_data}"
+                )
+            filter_data_.append(it)
+        elif isinstance(it, int):
+            if it < 1 or it > 27:
+                raise ValueError(
+                    f"The dataset with the ID={it} is not an "
+                    f"available dataset. The IDs are "
+                    f"{range(1, 28)}"
+                )
+            filter_data_.append(MAP_ID_NAME[it])
+        else:
+            raise ValueError(
+                f"The value in the tuple should be str or int."
+                f" Got {type(it)} instead."
+            )
+    return filter_data_
+
+
+def _load_dataset(dataset_name, zenodo_dir, download_if_missing, verbose):
+    """Load a single dataset from disk, downloading if necessary."""
+    filename = PRE_FILENAME + str(MAP_NAME_ID[dataset_name]) + POST_FILENAME
+    filepath = join(zenodo_dir, filename)
+    available = isfile(filepath)
+
+    if download_if_missing and not available:
+        makedirs(zenodo_dir, exist_ok=True)
+        if verbose:
+            print("Downloading %s" % URL)
+        f = BytesIO(urlopen(URL).read())
+        tar = tarfile.open(fileobj=f)
+        tar.extractall(path=zenodo_dir)
+    elif not download_if_missing and not available:
+        raise IOError("Data not found and `download_if_missing` is False")
+
+    data = np.load(filepath)
+    return data["data"], data["label"]
+
+
 @_deprecate_positional_args
 def fetch_zenodo_datasets(
     *,
@@ -165,11 +216,11 @@ def fetch_zenodo_datasets(
     +--+--------------+-------------------------------+-------+---------+-----+
     |2 |optical_digits| UCI, target: 8                | 9.1:1 | 5,620   | 64  |
     +--+--------------+-------------------------------+-------+---------+-----+
-    |3 |satimage      | UCI, target: 4                | 9.3:1 | 6,435   | 36  |
+    |3 |satimage      | UCI, target: 4                | 9.4:1 | 6,435   | 36  |
     +--+--------------+-------------------------------+-------+---------+-----+
     |4 |pen_digits    | UCI, target: 5                | 9.4:1 | 10,992  | 16  |
     +--+--------------+-------------------------------+-------+---------+-----+
-    |5 |abalone       | UCI, target: 7                | 9.7:1 | 4,177   | 10  |
+    |5 |abalone       | UCI, target: 7                | 9.4:1 | 4,177   | 10  |
     +--+--------------+-------------------------------+-------+---------+-----+
     |6 |sick_euthyroid| UCI, target: sick euthyroid   | 9.8:1 | 3,163   | 42  |
     +--+--------------+-------------------------------+-------+---------+-----+
@@ -227,55 +278,10 @@ def fetch_zenodo_datasets(
     zenodo_dir = join(data_home, "zenodo")
     datasets = OrderedDict()
 
-    if filter_data is None:
-        filter_data_ = MAP_NAME_ID.keys()
-    else:
-        list_data = MAP_NAME_ID.keys()
-        filter_data_ = []
-        for it in filter_data:
-            if isinstance(it, str):
-                if it not in list_data:
-                    raise ValueError(
-                        f"{it} is not a dataset available. "
-                        f"The available datasets are {list_data}"
-                    )
-                else:
-                    filter_data_.append(it)
-            elif isinstance(it, int):
-                if it < 1 or it > 27:
-                    raise ValueError(
-                        f"The dataset with the ID={it} is not an "
-                        f"available dataset. The IDs are "
-                        f"{range(1, 28)}"
-                    )
-                else:
-                    # The index start at one, then we need to remove one
-                    # to not have issue with the indexing.
-                    filter_data_.append(MAP_ID_NAME[it])
-            else:
-                raise ValueError(
-                    f"The value in the tuple should be str or int."
-                    f" Got {type(it)} instead."
-                )
+    filter_data_ = _resolve_filter_data(filter_data)
 
-    # go through the list and check if the data are available
     for it in filter_data_:
-        filename = PRE_FILENAME + str(MAP_NAME_ID[it]) + POST_FILENAME
-        filename = join(zenodo_dir, filename)
-        available = isfile(filename)
-
-        if download_if_missing and not available:
-            makedirs(zenodo_dir, exist_ok=True)
-            if verbose:
-                print("Downloading %s" % URL)
-            f = BytesIO(urlopen(URL).read())
-            tar = tarfile.open(fileobj=f)
-            tar.extractall(path=zenodo_dir)
-        elif not download_if_missing and not available:
-            raise IOError("Data not found and `download_if_missing` is False")
-
-        data = np.load(filename)
-        X, y = data["data"], data["label"]
+        X, y = _load_dataset(it, zenodo_dir, download_if_missing, verbose)
 
         if shuffle:
             ind = np.arange(X.shape[0])